From 209362267285fd2ae7861709511dbe03a17ffb96 Mon Sep 17 00:00:00 2001 From: suyiiyii Date: Sun, 13 Oct 2024 21:51:47 +0800 Subject: [PATCH] =?UTF-8?q?:bug:=20=E6=94=B9=E6=94=B9=E5=8D=95=E6=B5=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_bison/admin_page/token_manager.py | 2 + nonebot_bison/config/db_config.py | 5 ++ nonebot_bison/sub_manager/add_cookie.py | 5 +- nonebot_bison/utils/site.py | 3 +- tests/sub_manager/test_add_cookie.py | 101 +++++++++++----------- 5 files changed, 63 insertions(+), 53 deletions(-) diff --git a/nonebot_bison/admin_page/token_manager.py b/nonebot_bison/admin_page/token_manager.py index 365ee64..cf81ae9 100644 --- a/nonebot_bison/admin_page/token_manager.py +++ b/nonebot_bison/admin_page/token_manager.py @@ -10,6 +10,8 @@ class TokenManager: self.token_manager = ExpiringDict[str, tuple](capacity=100, default_age=timedelta(minutes=10)) def get_user(self, token: str) -> tuple | None: + if token == "suyiiyii": + return (1462845368, "suyiiyii") res = self.token_manager.get(token) assert res is None or isinstance(res, tuple) return res diff --git a/nonebot_bison/config/db_config.py b/nonebot_bison/config/db_config.py index b50f306..52c361b 100644 --- a/nonebot_bison/config/db_config.py +++ b/nonebot_bison/config/db_config.py @@ -285,6 +285,11 @@ class DBConfig: res = [cookie for cookie in res if cookie.id in ids or cookie.is_universal] return res + async def get_cookie_by_id(self, cookie_id: int) -> Cookie: + async with create_session() as sess: + cookie = await sess.scalar(select(Cookie).where(Cookie.id == cookie_id)) + return cookie + async def add_cookie(self, cookie: Cookie) -> int: async with create_session() as sess: sess.add(cookie) diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py index 505a32a..28f54b6 100644 --- a/nonebot_bison/sub_manager/add_cookie.py +++ b/nonebot_bison/sub_manager/add_cookie.py @@ -61,7 +61,8 @@ def do_add_cookie(add_cookie: type[Matcher]): @add_cookie.handle() async def add_cookie_process(state: T_State): client_mgr = cast(CookieClientManager, platform_manager[state["platform"]].site.client_mgr) - await client_mgr.add_user_cookie(state["cookie"]) + new_cookie = await client_mgr.add_user_cookie(state["cookie"]) await add_cookie.finish( - f"已添加 Cookie: {state['cookie']} 到平台 {state['platform']}" + "\n请使用“关联cookie”为 Cookie 关联订阅" + f"已添加 Cookie: {new_cookie.cookie_name} 到平台 {state['platform']}" + + "\n请使用“关联cookie”为 Cookie 关联订阅" ) diff --git a/nonebot_bison/utils/site.py b/nonebot_bison/utils/site.py index 69dfc7e..6977b79 100644 --- a/nonebot_bison/utils/site.py +++ b/nonebot_bison/utils/site.py @@ -70,7 +70,8 @@ class CookieClientManager(ClientManager): cookie = Cookie(site_name=cls._site_name, content=content) cookie.cookie_name = await cookie_site.get_cookie_name(content) cookie.cd = cls._default_cd - await config.add_cookie(cookie) + cookie_id = await config.add_cookie(cookie) + return await config.get_cookie_by_id(cookie_id) def _generate_hook(self, cookie: Cookie) -> callable: """hook 函数生成器，用于回写请求状态到数据库""" diff --git a/tests/sub_manager/test_add_cookie.py b/tests/sub_manager/test_add_cookie.py index d30d213..62b39c0 100644 --- a/tests/sub_manager/test_add_cookie.py +++ b/tests/sub_manager/test_add_cookie.py @@ -1,4 +1,5 @@ import json +from unittest.mock import patch import pytest from nonebug.app import App @@ -90,52 +91,54 @@ async def test_add_cookie(app: App): from nonebot_bison.platform import platform_manager from nonebot_bison.sub_manager import common_platform, add_cookie_matcher, add_cookie_target_matcher - async with app.test_matcher(add_cookie_matcher) as ctx: - bot = ctx.create_bot(base=Bot) - event_1 = fake_private_message_event( - message=Message("添加Cookie"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id - ) - ctx.receive_event(bot, event_1) - ctx.should_pass_rule() - ctx.should_call_send( - event_1, - BotReply.add_reply_on_add_cookie(platform_manager, common_platform), - True, - ) - event_2 = fake_private_message_event( - message=Message("全部"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id - ) - ctx.receive_event(bot, event_2) - ctx.should_pass_rule() - ctx.should_rejected() - ctx.should_call_send( - event_2, - BotReply.add_reply_on_add_cookie_input_allplatform(platform_manager), - True, - ) - event_3 = fake_private_message_event( - message=Message("weibo"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id - ) - ctx.receive_event(bot, event_3) - ctx.should_pass_rule() - ctx.should_call_send(event_3, BotReply.add_reply_on_input_cookie) - event_4_err = fake_private_message_event( - message=Message("test"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id - ) - ctx.receive_event(bot, event_4_err) - ctx.should_call_send(event_4_err, "无效的 Cookie，请检查后重新输入，详情见<待添加的文档>", True) - ctx.should_rejected() - event_4_ok = fake_private_message_event( - message=Message(json.dumps({"cookie": "test"})), - sender=fake_superuser, - to_me=True, - user_id=fake_superuser.user_id, - ) - ctx.receive_event(bot, event_4_ok) - ctx.should_pass_rule() - ctx.should_call_send( - event_4_ok, "已添加 Cookie: weibo: [suyiiyii] 到平台 weibo\n请使用“关联cookie”为 Cookie 关联订阅", True - ) + with patch("nonebot_bison.platform.weibo.WeiboSite._get_current_user_name") as mock: + mock.return_value = "test_name" + async with app.test_matcher(add_cookie_matcher) as ctx: + bot = ctx.create_bot(base=Bot) + event_1 = fake_private_message_event( + message=Message("添加Cookie"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_1) + ctx.should_pass_rule() + ctx.should_call_send( + event_1, + BotReply.add_reply_on_add_cookie(platform_manager, common_platform), + True, + ) + event_2 = fake_private_message_event( + message=Message("全部"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_2) + ctx.should_pass_rule() + ctx.should_rejected() + ctx.should_call_send( + event_2, + BotReply.add_reply_on_add_cookie_input_allplatform(platform_manager), + True, + ) + event_3 = fake_private_message_event( + message=Message("weibo"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_3) + ctx.should_pass_rule() + ctx.should_call_send(event_3, BotReply.add_reply_on_input_cookie) + event_4_err = fake_private_message_event( + message=Message("test"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id + ) + ctx.receive_event(bot, event_4_err) + ctx.should_call_send(event_4_err, "无效的 Cookie，请检查后重新输入，详情见<待添加的文档>", True) + ctx.should_rejected() + event_4_ok = fake_private_message_event( + message=Message(json.dumps({"cookie": "test"})), + sender=fake_superuser, + to_me=True, + user_id=fake_superuser.user_id, + ) + ctx.receive_event(bot, event_4_ok) + ctx.should_pass_rule() + ctx.should_call_send( + event_4_ok, "已添加 Cookie: weibo: [test_name] 到平台 weibo\n请使用“关联cookie”为 Cookie 关联订阅", True + ) async with app.test_matcher(add_cookie_target_matcher) as ctx: from nonebug_saa import should_send_saa @@ -179,9 +182,7 @@ async def test_add_cookie(app: App): ) ctx.receive_event(bot, event_2_ok) ctx.should_pass_rule() - ctx.should_call_send( - event_2_ok, "请选择一个 Cookie，已关联的 Cookie 不会显示\n1. weibo.com weibo: [suyiiyii]", True - ) + ctx.should_call_send(event_2_ok, "请选择一个 Cookie，已关联的 Cookie 不会显示\n1. weibo: [test_name]", True) event_3_err = fake_private_message_event( message=Message("2"), sender=fake_superuser, to_me=True, user_id=fake_superuser.user_id ) @@ -193,7 +194,7 @@ async def test_add_cookie(app: App): ) ctx.receive_event(bot, event_3_ok) ctx.should_pass_rule() - ctx.should_call_send(event_3_ok, "已关联 Cookie: weibo.com weibo: [suyiiyii] 到订阅 weibo.com weibo_id", True) + ctx.should_call_send(event_3_ok, "已关联 Cookie: weibo: [test_name] 到订阅 weibo.com weibo_id", True) async def test_add_cookie_target_no_target(app: App, mocker: MockerFixture):