Refactor BaseClient class and add cookie property accessor

Added a `cookies` property accessor to the `BaseClient` class, which retrieves cookies from the `httpx` AsyncClient object and returns a `simnet.client.cookies.Cookies` object (which is a custom implementation of the `httpx.Cookies` class). Additionally, the initialization function was refactored to better support using cookies. Now, if `account_id` is not passed in, it will be retrieved from the cookies.
This commit is contained in:
洛水居室 2023-05-09 15:19:36 +08:00 committed by GitHub
parent 5ed80f626e
commit b74b9aef33
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 33 additions and 18 deletions

View File

@ -38,13 +38,13 @@ class BaseClient(AsyncContextManager["BaseClient"]):
timeout (Optional[TimeoutTypes], optional): Timeout configuration for the client.
Attributes:
cookies (CookieTypes): The cookies used for the client.
headers (HeaderTypes): The headers used for the client.
account_id (Optional[int]): The account id used for the client.
player_id (Optional[int]): The player id used for the client.
region (Region): The region used for the client.
lang (str): The language used for the client.
game (Optional[Game]): The game used for the client.
"""
game: Optional[Game] = None
@ -69,18 +69,22 @@ class BaseClient(AsyncContextManager["BaseClient"]):
pool=1.0,
)
self.cookies = Cookies(cookies)
cookies = Cookies(cookies)
self.headers = Headers(headers)
self.player_id = player_id
self.account_id = account_id
self.client = AsyncClient(cookies=self.cookies, timeout=timeout)
self.account_id = account_id or cookies.account_id
self.client = AsyncClient(cookies=cookies, timeout=timeout)
self.region = region
self.lang = lang
def get_player_id(self) -> Optional[int]:
"""Get the player id used for the client."""
player_id = self.player_id or self.cookies.account_id
return player_id
@property
def cookies(self) -> Cookies:
"""Get the cookies used for the client."""
return Cookies(self.client.cookies.jar)
@cookies.setter
def cookies(self, cookies: CookieTypes) -> None:
self.client.cookies = cookies
@property
def device_name(self) -> str:
@ -202,9 +206,7 @@ class BaseClient(AsyncContextManager["BaseClient"]):
if self.region == Region.OVERSEAS:
header["x-rpc-language"] = self.lang or lang
if ds is None:
app_version, client_type, ds = generate_dynamic_secret(
self.region, ds_type, new_ds, data, params
)
app_version, client_type, ds = generate_dynamic_secret(self.region, ds_type, new_ds, data, params)
header["x-rpc-app_version"] = app_version
header["x-rpc-client_type"] = client_type
header["DS"] = ds
@ -333,9 +335,5 @@ class BaseClient(AsyncContextManager["BaseClient"]):
"""
if method is None:
method = "POST" if data else "GET"
headers = self.get_lab_api_header(
headers, ds_type=ds_type, new_ds=new_ds, lang=lang, data=data, params=params
)
return await self.request_api(
method=method, url=url, json=data, params=params, headers=headers
)
headers = self.get_lab_api_header(headers, ds_type=ds_type, new_ds=new_ds, lang=lang, data=data, params=params)
return await self.request_api(method=method, url=url, json=data, params=params, headers=headers)

View File

@ -8,7 +8,7 @@ __all__ = ("Cookies",)
class Cookies(_Cookies):
"""An extension of the `httpx.Cookies` class that includes additional functionality."""
COOKIE_USER_ID_NAMES = ("ltuid", "account_id", "ltuid_v2", "account_id_v2")
COOKIE_USER_ID_NAMES = ("ltuid", "account_id", "stuid", "ltuid_v2", "account_id_v2")
@property
def account_id(self) -> Optional[int]:

17
tests/test_base_client.py Normal file
View File

@ -0,0 +1,17 @@
import pytest
from simnet.client.base import BaseClient
from simnet.client.cookies import Cookies
@pytest.mark.asyncio
class TestBaseClient:
@staticmethod
async def test_cookies():
async with BaseClient(cookies={"uid": "114514"}) as client:
assert isinstance(client.cookies, Cookies)
client.cookies = {"account_id": "114514"}
assert isinstance(client.cookies, Cookies)
assert client.cookies.get("account_id") == "114514"
client.cookies.set("stuid", "114514")
assert client.cookies.get("stuid") == "114514"