From 29c2eb456dd7849da613b488de9495c84e830ad9 Mon Sep 17 00:00:00 2001
From: suyiiyii <suyiiyii@gmail.com>
Date: Tue, 29 Oct 2024 12:24:31 +0800
Subject: [PATCH] =?UTF-8?q?:recycle:=20=E6=9B=B4=E6=96=B0=E9=83=A8?=
=?UTF-8?q?=E5=88=86=E6=96=B9=E6=B3=95=E4=BD=BF=E7=94=A8=20scheduler=5Fdic?=
=?UTF-8?q?t=20=E8=BF=9B=E8=A1=8Ccookie=E6=93=8D=E4=BD=9C?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
---
nonebot_bison/admin_page/api.py | 3 ++-
nonebot_bison/sub_manager/add_cookie.py | 5 +++--
tests/config/test_cookie.py | 3 ++-
tests/platforms/test_bilibili.py | 2 +-
4 files changed, 8 insertions(+), 5 deletions(-)
diff --git a/nonebot_bison/admin_page/api.py b/nonebot_bison/admin_page/api.py
index fc3a31f..b5d90cc 100644
--- a/nonebot_bison/admin_page/api.py
+++ b/nonebot_bison/admin_page/api.py
@@ -12,6 +12,7 @@ from fastapi.security.oauth2 import OAuth2PasswordBearer
from ..types import WeightConfig
from ..apis import check_sub_target
from .jwt import load_jwt, pack_jwt
+from ..scheduler import scheduler_dict
from ..types import Target as T_Target
from ..utils.get_bot import get_groups
from .token_manager import token_manager
@@ -231,7 +232,7 @@ async def get_cookie(site_name: str = None, target: str = None) -> list[Cookie]:
@router.post("/cookie", dependencies=[Depends(check_is_superuser)])
async def add_cookie(site_name: str, content: str) -> StatusResp:
- client_mgr = cast(CookieClientManager, site_manager[site_name].client_mgr)
+ client_mgr = cast(CookieClientManager, scheduler_dict[site_manager[site_name]].client_mgr)
await client_mgr.add_user_cookie(content)
return StatusResp(ok=True, msg="")
diff --git a/nonebot_bison/sub_manager/add_cookie.py b/nonebot_bison/sub_manager/add_cookie.py
index 65ff11b..c22163a 100644
--- a/nonebot_bison/sub_manager/add_cookie.py
+++ b/nonebot_bison/sub_manager/add_cookie.py
@@ -8,6 +8,7 @@ from nonebot.params import Arg, ArgPlainText
from nonebot.adapters.onebot.v11 import MessageEvent
from nonebot.adapters import Message, MessageTemplate
+from ..scheduler import scheduler_dict
from ..platform import platform_manager
from .utils import common_platform, gen_handle_cancel, only_allow_private
from ..utils.site import CookieSite, CookieClientManager, is_cookie_client_manager
@@ -46,7 +47,6 @@ def do_add_cookie(add_cookie: type[Matcher]):
await add_cookie.finish("已中止添加cookie")
elif platform in platform_manager:
state["platform"] = platform
- state["site"] = cast(CookieSite, platform_manager[platform].site)
else:
await add_cookie.reject("平台输入错误")
@@ -74,7 +74,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)
+ cookie_site = cast(type[CookieSite], platform_manager[state["platform"]].site)
+ client_mgr = cast(CookieClientManager, scheduler_dict[cookie_site].client_mgr)
new_cookie = await client_mgr.add_user_cookie(state["cookie"], state["cookie_name"])
await add_cookie.finish(
f"已添加 Cookie: {new_cookie.cookie_name} 到平台 {state['platform']}"
diff --git a/tests/config/test_cookie.py b/tests/config/test_cookie.py
index 549cb7f..9f29790 100644
--- a/tests/config/test_cookie.py
+++ b/tests/config/test_cookie.py
@@ -12,6 +12,7 @@ async def test_cookie(app: App, init_scheduler):
from nonebot_bison.platform import site_manager
from nonebot_bison.config.db_config import config
+ from nonebot_bison.scheduler import scheduler_dict
from nonebot_bison.types import Target as T_Target
from nonebot_bison.utils.site import CookieClientManager
from nonebot_bison.config.utils import DuplicateCookieTargetException
@@ -27,7 +28,7 @@ async def test_cookie(app: App, init_scheduler):
tags=[],
)
site = site_manager["weibo.com"]
- client_mgr = cast(CookieClientManager, site.client_mgr)
+ client_mgr = cast(CookieClientManager, scheduler_dict[site].client_mgr)
# 刷新匿名cookie
await client_mgr.refresh_client()
diff --git a/tests/platforms/test_bilibili.py b/tests/platforms/test_bilibili.py
index e31dfd8..95e1846 100644
--- a/tests/platforms/test_bilibili.py
+++ b/tests/platforms/test_bilibili.py
@@ -32,7 +32,7 @@ def bilibili(app: App) -> "Bilibili":
from nonebot_bison.platform import platform_manager
from nonebot_bison.platform.bilibili import BilibiliClientManager
- return platform_manager["bilibili"](ProcessContext(BilibiliClientManager())) # type: ignore
+ return platform_manager["bilibili"](ProcessContext(BilibiliClientManager(platform_manager["bilibili"].site))) # type: ignore
@pytest.fixture