fix: support ugreen token_id login response

tests/test_ugreen_api.py

@@ -1,21 +1,28 @@
-import unittest
 from types import SimpleNamespace
+from typing import Optional
 from unittest.mock import patch
 
 from app.modules.ugreen.api import Api
 
 
 class _FakeResponse:
-    def __init__(self, payload: dict, headers: dict | None = None):
+    def __init__(self, payload: dict, headers: Optional[dict] = None) -> None:
+        """初始化伪造 HTTP 响应"""
         self._payload = payload
         self.headers = headers or {}
 
-    def json(self):
+    def json(self) -> dict:
+        """返回伪造 JSON 响应体"""
         return self._payload
 
 
 class _FakeSession:
-    def __init__(self, get_responses=None, post_responses=None):
+    def __init__(
+        self,
+        get_responses: Optional[list[_FakeResponse]] = None,
+        post_responses: Optional[list[_FakeResponse]] = None,
+    ) -> None:
+        """初始化伪造 HTTP 会话"""
         self._get_responses = list(get_responses or [])
         self._post_responses = list(post_responses or [])
         self.calls: list[tuple[str, dict]] = []
@@ -24,86 +31,149 @@ class _FakeSession:
             update=lambda *_args, **_kwargs: None,
         )
 
-    def get(self, *args, **kwargs):
+    def get(self, *args: object, **kwargs: object) -> _FakeResponse:
+        """记录 GET 请求并返回预置响应"""
         if args:
             kwargs = {"url": args[0], **kwargs}
         self.calls.append(("GET", kwargs))
         return self._get_responses.pop(0) if self._get_responses else _FakeResponse({})
 
-    def post(self, *args, **kwargs):
+    def post(self, *args: object, **kwargs: object) -> _FakeResponse:
+        """记录 POST 请求并返回预置响应"""
         if args:
             kwargs = {"url": args[0], **kwargs}
         self.calls.append(("POST", kwargs))
         return self._post_responses.pop(0) if self._post_responses else _FakeResponse({})
 
     @staticmethod
-    def close():
+    def close() -> None:
+        """关闭伪造会话"""
         return None
 
 
 class _FakeCrypto:
-    def __init__(self, *args, **kwargs):
-        pass
+    def __init__(self, *args: object, **kwargs: object) -> None:
+        """初始化伪造加密工具"""
+        return None
 
     @staticmethod
     def rsa_encrypt_long(raw: str) -> str:
+        """返回可断言的伪造 RSA 加密内容"""
         return f"enc:{raw}"
 
     @staticmethod
-    def build_encrypted_request(url: str, method: str = "GET", params=None, **kwargs):
-        return SimpleNamespace(url=url, headers={}, params=params or {}, json=None, aes_key="k")
+    def build_encrypted_request(
+        url: str,
+        method: str = "GET",
+        params: Optional[dict] = None,
+        **kwargs: object,
+    ) -> SimpleNamespace:
+        """构造伪造加密请求对象"""
+        _ = method, kwargs
+        return SimpleNamespace(
+            url=url,
+            headers={},
+            params=params or {},
+            json=None,
+            aes_key="k",
+        )
 
     @staticmethod
-    def decrypt_response(payload, aes_key):
+    def decrypt_response(payload: dict, aes_key: str) -> dict:
+        """原样返回伪造响应内容"""
+        _ = aes_key
         return payload
 
 
-class UgreenApiVerifySslTest(unittest.TestCase):
-    def test_request_json_default_verify_ssl_true(self):
-        api = Api(host="https://example.com")
-        fake_session = _FakeSession(
-            get_responses=[_FakeResponse({"code": 200})],
-            post_responses=[_FakeResponse({"code": 200})],
-        )
-        api._session = fake_session
+def test_request_json_default_verify_ssl_true() -> None:
+    """默认请求应开启 HTTPS 证书校验"""
+    api = Api(host="https://example.com")
+    fake_session = _FakeSession(
+        get_responses=[_FakeResponse({"code": 200})],
+        post_responses=[_FakeResponse({"code": 200})],
+    )
+    api._session = fake_session
 
-        api._request_json(url="https://example.com/a", method="GET")
-        api._request_json(url="https://example.com/b", method="POST", json_data={"x": 1})
+    api._request_json(url="https://example.com/a", method="GET")
+    api._request_json(url="https://example.com/b", method="POST", json_data={"x": 1})
 
-        self.assertEqual(fake_session.calls[0][1].get("verify"), True)
-        self.assertEqual(fake_session.calls[1][1].get("verify"), True)
+    assert fake_session.calls[0][1].get("verify") is True
+    assert fake_session.calls[1][1].get("verify") is True
 
-    def test_login_logout_follow_verify_ssl_flag(self):
-        api = Api(host="https://example.com", verify_ssl=False)
-        fake_session = _FakeSession(
-            get_responses=[_FakeResponse({})],
-            post_responses=[
-                _FakeResponse({"code": 200, "msg": "ok", "data": {}}, headers={"x-rsa-token": "BEGIN TEST"}),
-                _FakeResponse(
-                    {
-                        "code": 200,
-                        "msg": "ok",
-                        "data": {
-                            "token": "token-value",
-                            "public_key": "BEGIN LOGIN KEY",
-                            "static_token": "static-token",
-                            "is_ugk": False,
-                        },
-                    }
-                ),
-            ],
-        )
-        api._session = fake_session
 
-        with patch("app.modules.ugreen.api.UgreenCrypto", _FakeCrypto):
-            token = api.login("tester", "pwd")
-            self.assertEqual(token, "token-value")
-            api.logout()
+def test_login_logout_follow_verify_ssl_flag() -> None:
+    """登录与登出请求应沿用用户配置的证书校验开关"""
+    api = Api(host="https://example.com", verify_ssl=False)
+    fake_session = _FakeSession(
+        get_responses=[_FakeResponse({})],
+        post_responses=[
+            _FakeResponse(
+                {"code": 200, "msg": "ok", "data": {}},
+                headers={"x-rsa-token": "BEGIN CHECK KEY"},
+            ),
+            _FakeResponse(
+                {
+                    "code": 200,
+                    "msg": "ok",
+                    "data": {
+                        "token": "token-value",
+                        "public_key": "BEGIN LOGIN KEY",
+                        "static_token": "static-token",
+                        "is_ugk": False,
+                    },
+                }
+            ),
+        ],
+    )
+    api._session = fake_session
 
-        self.assertEqual(len(fake_session.calls), 3)
-        self.assertEqual(fake_session.calls[0][0], "POST")
-        self.assertEqual(fake_session.calls[1][0], "POST")
-        self.assertEqual(fake_session.calls[2][0], "GET")
-        self.assertEqual(fake_session.calls[0][1].get("verify"), False)
-        self.assertEqual(fake_session.calls[1][1].get("verify"), False)
-        self.assertEqual(fake_session.calls[2][1].get("verify"), False)
+    with patch("app.modules.ugreen.api.UgreenCrypto", _FakeCrypto):
+        token = api.login("tester", "pwd")
+        assert token == "token-value"
+        assert api.public_key == "BEGIN LOGIN KEY"
+        assert api.static_token == "static-token"
+        api.logout()
+
+    assert len(fake_session.calls) == 3
+    assert fake_session.calls[0][0] == "POST"
+    assert fake_session.calls[1][0] == "POST"
+    assert fake_session.calls[2][0] == "GET"
+    assert fake_session.calls[0][1].get("verify") is False
+    assert fake_session.calls[1][1].get("verify") is False
+    assert fake_session.calls[2][1].get("verify") is False
+
+
+def test_login_accepts_token_id_and_reuses_check_public_key() -> None:
+    """登录响应只有 token_id 时应复用检查接口公钥完成会话初始化"""
+    api = Api(host="https://example.com")
+    fake_session = _FakeSession(
+        post_responses=[
+            _FakeResponse(
+                {"code": 200, "msg": "ok", "data": {}},
+                headers={"x-rsa-token": "BEGIN CHECK KEY"},
+            ),
+            _FakeResponse(
+                {
+                    "code": 200,
+                    "msg": "success",
+                    "data": {
+                        "enable_otp": True,
+                        "is_exceed": False,
+                        "role": "admin",
+                        "token_id": "token-id-value",
+                        "uid": 1000,
+                        "urgent_email": "tester@example.com",
+                    },
+                }
+            ),
+        ],
+    )
+    api._session = fake_session
+
+    with patch("app.modules.ugreen.api.UgreenCrypto", _FakeCrypto):
+        token = api.login("tester", "pwd")
+
+    assert token == "token-id-value"
+    assert api.token == "token-id-value"
+    assert api.static_token == "token-id-value"
+    assert api.public_key == "BEGIN CHECK KEY"