From b74b9aef33383d5f3eeadf97cbcbe60dc4ece11d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B4=9B=E6=B0=B4=E5=B1=85=E5=AE=A4?=
 <luoshuijs@outlook.com>
Date: Tue, 9 May 2023 15:19:36 +0800
Subject: [PATCH] :sparkles: 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.
---
 simnet/client/base.py     | 32 +++++++++++++++-----------------
 simnet/client/cookies.py  |  2 +-
 tests/test_base_client.py | 17 +++++++++++++++++
 3 files changed, 33 insertions(+), 18 deletions(-)
 create mode 100644 tests/test_base_client.py

diff --git a/simnet/client/base.py b/simnet/client/base.py
index 722473f..6b8b4c6 100644
--- a/simnet/client/base.py
+++ b/simnet/client/base.py
@@ -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)
diff --git a/simnet/client/cookies.py b/simnet/client/cookies.py
index b70aca7..e8cadba 100644
--- a/simnet/client/cookies.py
+++ b/simnet/client/cookies.py
@@ -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]:
diff --git a/tests/test_base_client.py b/tests/test_base_client.py
new file mode 100644
index 0000000..76761ec
--- /dev/null
+++ b/tests/test_base_client.py
@@ -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"