From 2cfd58373f5e4ef99b58bab389a85371da81eb19 Mon Sep 17 00:00:00 2001
From: suyiiyii <suyiiyii@gmail.com>
Date: Fri, 13 Sep 2024 14:12:02 +0800
Subject: [PATCH] :white_check_mark: pytest db_config

---
 nonebot_bison/config/db_config.py |  8 +++--
 tests/config/test_cookie.py       | 51 ++++++++++++++++++++++++++++++-
 2 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py
index 2f85775..9dac6dd 100644
--- a/nonebot_bison/config/db_config.py
+++ b/nonebot_bison/config/db_config.py
@@ -296,7 +296,7 @@ class DBConfig:
         async with create_session() as sess:
             cookie_in_db: Cookie | None = await sess.scalar(select(Cookie).where(Cookie.id == cookie.id))
             if not cookie_in_db:
-                return
+                raise ValueError(f"cookie {cookie.id} not found")
             cookie_in_db.content = cookie.content
             cookie_in_db.last_usage = cookie.last_usage
             cookie_in_db.status = cookie.status
@@ -333,9 +333,11 @@ class DBConfig:
             sess.add(cookie_target)
             await sess.commit()
 
-    async def delete_cookie_target(self, target: T_Target, site_name: str, cookie_id: int):
+    async def delete_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int):
         async with create_session() as sess:
-            target_obj = await sess.scalar(select(Target).where(Target.site_name == site_name, Target.target == target))
+            target_obj = await sess.scalar(
+                select(Target).where(Target.platform_name == platform_name, Target.target == target)
+            )
             cookie_obj = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id))
             await sess.execute(
                 delete(CookieTarget).where(CookieTarget.target == target_obj, CookieTarget.cookie == cookie_obj)
diff --git a/tests/config/test_cookie.py b/tests/config/test_cookie.py
index 7db0b47..c831176 100644
--- a/tests/config/test_cookie.py
+++ b/tests/config/test_cookie.py
@@ -1,6 +1,8 @@
 import json
 from typing import cast
+from datetime import datetime
 
+import pytest
 from nonebug import App
 
 
@@ -11,6 +13,7 @@ async def test_cookie(app: App, init_scheduler):
     from nonebot_bison.config.db_config import config
     from nonebot_bison.types import Target as T_Target
     from nonebot_bison.utils.site import CookieClientManager
+    from nonebot_bison.config.utils import DuplicateCookieTargetException
 
     target = T_Target("weibo_id")
     platform_name = "weibo"
@@ -25,11 +28,13 @@ async def test_cookie(app: App, init_scheduler):
     site = site_manager["weibo.com"]
     client_mgr = cast(CookieClientManager, site.client_mgr)
 
-    await client_mgr.refresh_anonymous_cookie()  # 刷新匿名cookie
+    # 刷新匿名cookie
+    await client_mgr.refresh_anonymous_cookie()
 
     cookies = await config.get_cookie(site_name=site.name)
     assert len(cookies) == 1
 
+    # 添加用户cookie
     await client_mgr.add_user_cookie(json.dumps({"test_cookie": "1"}))
     await client_mgr.add_user_cookie(json.dumps({"test_cookie": "2"}))
 
@@ -39,6 +44,7 @@ async def test_cookie(app: App, init_scheduler):
     cookies = await config.get_cookie(site_name=site.name, is_anonymous=False)
     assert len(cookies) == 2
 
+    # 单个target，多个cookie
     await config.add_cookie_target(target, platform_name, cookies[0].id)
     await config.add_cookie_target(target, platform_name, cookies[1].id)
 
@@ -64,11 +70,54 @@ async def test_cookie(app: App, init_scheduler):
     await client_mgr.add_user_cookie(json.dumps({"test_cookie": "3"}))
     cookies = await config.get_cookie(site_name=site.name, is_anonymous=False)
 
+    # 多个target，多个cookie
     await config.add_cookie_target(target2, platform_name, cookies[0].id)
     await config.add_cookie_target(target2, platform_name, cookies[2].id)
 
     cookies = await config.get_cookie(site_name=site.name, target=target2)
     assert len(cookies) == 3
 
+    # 重复关联 target
+    with pytest.raises(DuplicateCookieTargetException) as e:
+        await config.add_cookie_target(target2, platform_name, cookies[2].id)
+    assert isinstance(e.value, DuplicateCookieTargetException)
+
     cookies = await config.get_cookie(site_name=site.name, target=target2, is_anonymous=False)
     assert len(cookies) == 2
+
+    # 有关联的cookie不能删除
+    with pytest.raises(Exception, match="cookie") as e:
+        await config.delete_cookie_by_id(cookies[1].id)
+    cookies = await config.get_cookie(site_name=site.name, target=target2, is_anonymous=False)
+    assert len(cookies) == 2
+
+    await config.delete_cookie_target(target2, platform_name, cookies[1].id)
+    await config.delete_cookie_by_id(cookies[1].id)
+    cookies = await config.get_cookie(site_name=site.name, target=target2, is_anonymous=False)
+    assert len(cookies) == 1
+
+    cookie = cookies[0]
+    cookie_id = cookie.id
+    cookie.last_usage = datetime(2024, 9, 13)
+    cookie.status = "test"
+    await config.update_cookie(cookie)
+    cookies = await config.get_cookie(site_name=site.name, target=target2, is_anonymous=False)
+    assert len(cookies) == 1
+    assert cookies[0].id == cookie_id
+    assert cookies[0].last_usage == datetime(2024, 9, 13)
+    assert cookies[0].status == "test"
+
+    # 不存在的 cookie_id
+    cookie.id = 114514
+    with pytest.raises(ValueError, match="cookie") as e:
+        await config.update_cookie(cookie)
+
+    # 获取所有关联对象
+    cookie_targets = await config.get_cookie_target()
+    assert len(cookie_targets) == 3
+
+    # 删除关联对象
+    await config.delete_cookie_target_by_id(cookie_targets[0].id)
+
+    cookie_targets = await config.get_cookie_target()
+    assert len(cookie_targets) == 2