feat: support ApplicationOAuth jwt auth to impl Auth (#12)

cozepy/__init__.py

@@ -1,4 +1,4 @@
-from .auth import ApplicationOAuth, Auth, TokenAuth
+from .auth import ApplicationOAuth, Auth, TokenAuth, JWTAuth
 from .config import COZE_COM_BASE_URL, COZE_CN_BASE_URL
 from .coze import Coze
 from .model import (
@@ -16,6 +16,7 @@
     "ApplicationOAuth",
     "Auth",
     "TokenAuth",
+    "JWTAuth",
     "COZE_COM_BASE_URL",
     "COZE_CN_BASE_URL",
     "Coze",

cozepy/auth.py

@@ -56,7 +56,7 @@ def __init__(self, client_id: str, client_secret: str = "", base_url: str = COZE
         self._token = ""
         self._requester = Requester()
 
-    def jwt_auth(self, private_key: str, kid: str, ttl: int):
+    def jwt_auth(self, private_key: str, kid: str, ttl: int) -> OAuthToken:
         """
         Get the token by jwt with jwt auth flow.
         """
@@ -137,3 +137,40 @@ def token_type(self) -> str:
     @property
     def token(self) -> str:
         return self._token
+
+
+class JWTAuth(Auth):
+    """
+    The JWT auth flow.
+    """
+
+    def __init__(self, client_id: str, private_key: str, kid: str, ttl: int = 7200, base_url: str = COZE_COM_BASE_URL):
+        assert isinstance(client_id, str)
+        assert isinstance(private_key, str)
+        assert isinstance(kid, str)
+        assert isinstance(ttl, int)
+        assert ttl > 0
+        assert isinstance(base_url, str)
+
+        self._client_id = client_id
+        self._private_key = private_key
+        self._kid = kid
+        self._ttl = ttl
+        self._base_url = base_url
+        self._token = None
+        self._oauth_cli = ApplicationOAuth(self._client_id, base_url=self._base_url)
+
+    @property
+    def token_type(self) -> str:
+        return "Bearer"
+
+    @property
+    def token(self) -> str:
+        token = self._generate_token()
+        return token.access_token
+
+    def _generate_token(self):
+        if self._token is not None and int(time.time()) < self._token.expires_in:
+            return self._token
+        self._token = self._oauth_cli.jwt_auth(self._private_key, self._kid, self._ttl)
+        return self._token

cozepy/request.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Tuple, Optional, Union, List, get_origin, get_args
+from typing import TYPE_CHECKING, Tuple, Optional, Union, List, get_origin, get_args, Iterator
 
 import requests
 from requests import Response
@@ -37,7 +37,7 @@ def request(
         headers: dict = None,
         body: dict = None,
         stream: bool = False,
-    ) -> Union[T, List[T]]:
+    ) -> Union[T, List[T], Iterator[bytes]]:
         """
         Send a request to the server.
         """

tests/config.py

@@ -0,0 +1,20 @@
+import os
+
+from cozepy import ApplicationOAuth, COZE_CN_BASE_URL, JWTAuth, TokenAuth
+
+COZE_JWT_AUTH_CLIENT_ID = os.getenv("COZE_JWT_AUTH_CLIENT_ID").strip()
+COZE_JWT_AUTH_PRIVATE_KEY = os.getenv("COZE_JWT_AUTH_PRIVATE_KEY").strip()
+COZE_JWT_AUTH_KEY_ID = os.getenv("COZE_JWT_AUTH_KEY_ID").strip()
+
+COZE_TOKEN = os.getenv("COZE_TOKEN").strip()
+
+app_oauth = ApplicationOAuth(
+    COZE_JWT_AUTH_CLIENT_ID,
+    base_url=COZE_CN_BASE_URL,
+)
+
+fixed_token_auth = TokenAuth(COZE_TOKEN)
+
+jwt_auth = JWTAuth(
+    COZE_JWT_AUTH_CLIENT_ID, COZE_JWT_AUTH_PRIVATE_KEY, COZE_JWT_AUTH_KEY_ID, 30, base_url=COZE_CN_BASE_URL
+)

tests/test_auth.py

@@ -1,20 +1,15 @@
-import os
 import time
 
-from cozepy import ApplicationOAuth, COZE_CN_BASE_URL
+from tests.config import app_oauth, COZE_JWT_AUTH_KEY_ID, COZE_JWT_AUTH_PRIVATE_KEY, jwt_auth
 
 
-def test_jwt_auth():
-    client_id = os.getenv("COZE_JWT_AUTH_CLIENT_ID")
-    private_key = os.getenv("COZE_JWT_AUTH_PRIVATE_KEY")
-    key_id = os.getenv("COZE_JWT_AUTH_KEY_ID")
-
-    app = ApplicationOAuth(
-        client_id,
-        base_url=COZE_CN_BASE_URL,
-    )
-    token = app.jwt_auth(private_key, key_id, 30)
+def test_jwt_app_oauth():
+    token = app_oauth.jwt_auth(COZE_JWT_AUTH_PRIVATE_KEY, COZE_JWT_AUTH_KEY_ID, 30)
     assert token.access_token != ""
     assert token.token_type == "Bearer"
     assert token.expires_in - int(time.time()) <= 31
     assert token.refresh_token == ""
+
+
+def test_jwt_auth():
+    assert jwt_auth.token != ""

tests/test_bot.py

@@ -1,17 +1,45 @@
 import os
 from unittest import TestCase
 
-from cozepy import TokenAuth, Coze, COZE_CN_BASE_URL
+from cozepy import Coze, COZE_CN_BASE_URL
+from tests.config import fixed_token_auth, jwt_auth
 
 
 class TestBotClient(TestCase):
-    def test_list_published_bots_v1(self):
+    def test_bot_v1_list(self):
         space_id = os.getenv("SPACE_ID_1").strip()
-        token = os.getenv("COZE_TOKEN").strip()
-        auth = TokenAuth(token)
-        cli = Coze(auth=auth, base_url=COZE_CN_BASE_URL)
 
+        cli_list = [
+            # fixed token
+            Coze(auth=fixed_token_auth, base_url=COZE_CN_BASE_URL),
+            # jwt auth
+            Coze(auth=jwt_auth, base_url=COZE_CN_BASE_URL),
+        ]
+        for cli in cli_list:
+            res = cli.bot.v1.list(space_id=space_id, page_size=2)
+            assert res.total > 1
+            assert res.has_more
+            assert len(res.items) > 1
+
+    def test_bot_v1_get_online_info(self):
+        bot_id = self.bot_id
+
+        cli_list = [
+            # fixed token
+            Coze(auth=fixed_token_auth, base_url=COZE_CN_BASE_URL),
+            # jwt auth
+            Coze(auth=jwt_auth, base_url=COZE_CN_BASE_URL),
+        ]
+        for cli in cli_list:
+            bot = cli.bot.v1.get_online_info(bot_id=bot_id)
+            assert bot is not None
+            assert bot.bot_id == bot_id
+
+    @property
+    def bot_id(self) -> str:
+        space_id = os.getenv("SPACE_ID_1").strip()
+
+        # fixed token
+        cli = Coze(auth=fixed_token_auth, base_url=COZE_CN_BASE_URL)
         res = cli.bot.v1.list(space_id=space_id, page_size=2)
-        assert res.total > 1
-        assert res.has_more
-        assert len(res.items) > 1
+        return res.items[0].bot_id