pytest db_config

This commit is contained in:
suyiiyii 2024-09-13 14:12:02 +08:00
parent 4791fb69e0
commit 2cfd58373f
2 changed files with 55 additions and 4 deletions

View File

@ -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)

View File

@ -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