feat: restructure and extract common class and func

Author: SoulPancake

openfga_sdk/oauth2.py

@@ -1,10 +1,7 @@
 import asyncio
 import json
-import math
 import random
-import sys
 
-from dataclasses import dataclass
 from datetime import datetime, timedelta
 
 import urllib3
@@ -17,32 +14,11 @@
 )
 from openfga_sdk.credentials import Credentials
 from openfga_sdk.exceptions import AuthenticationError
+from openfga_sdk.oauth2_common import _TokenState, jitter
 from openfga_sdk.telemetry.attributes import TelemetryAttributes
 from openfga_sdk.telemetry.telemetry import Telemetry
 
 
-@dataclass(frozen=True)
-class _TokenState:
-    access_token: str
-    expiry_time: datetime
-    expiry_buffer: float
-
-
-def jitter(loop_count, min_wait_in_ms):
-    """
-    Generate a random jitter value for exponential backoff
-    """
-    minimum = math.ceil(2**loop_count * min_wait_in_ms)
-    maximum = math.ceil(2 ** (loop_count + 1) * min_wait_in_ms)
-    jitter = random.randrange(minimum, maximum) / 1000
-
-    # If running in pytest, set jitter to 0 to speed up tests
-    if "pytest" in sys.modules:
-        jitter = 0
-
-    return jitter
-
-
 class OAuth2Client:
     def __init__(self, credentials: Credentials, configuration=None):
         self._credentials = credentials

openfga_sdk/oauth2_common.py

@@ -0,0 +1,28 @@
+import math
+import random
+import sys
+
+from dataclasses import dataclass
+from datetime import datetime
+
+
+@dataclass(frozen=True)
+class _TokenState:
+    access_token: str
+    expiry_time: datetime
+    expiry_buffer: float
+
+
+def jitter(loop_count, min_wait_in_ms):
+    """
+    Generate a random jitter value for exponential backoff
+    """
+    minimum = math.ceil(2**loop_count * min_wait_in_ms)
+    maximum = math.ceil(2 ** (loop_count + 1) * min_wait_in_ms)
+    jitter = random.randrange(minimum, maximum) / 1000
+
+    # If running in pytest, set jitter to 0 to speed up tests
+    if "pytest" in sys.modules:
+        jitter = 0
+
+    return jitter

openfga_sdk/sync/oauth2.py

@@ -1,11 +1,8 @@
 import json
-import math
 import random
-import sys
 import threading
 import time
 
-from dataclasses import dataclass
 from datetime import datetime, timedelta
 
 import urllib3
@@ -18,32 +15,11 @@
 )
 from openfga_sdk.credentials import Credentials
 from openfga_sdk.exceptions import AuthenticationError
+from openfga_sdk.oauth2_common import _TokenState, jitter
 from openfga_sdk.telemetry.attributes import TelemetryAttributes
 from openfga_sdk.telemetry.telemetry import Telemetry
 
 
-@dataclass(frozen=True)
-class _TokenState:
-    access_token: str
-    expiry_time: datetime
-    expiry_buffer: float
-
-
-def jitter(loop_count, min_wait_in_ms):
-    """
-    Generate a random jitter value for exponential backoff
-    """
-    minimum = math.ceil(2**loop_count * min_wait_in_ms)
-    maximum = math.ceil(2 ** (loop_count + 1) * min_wait_in_ms)
-    jitter = random.randrange(minimum, maximum) / 1000
-
-    # If running in pytest, set jitter to 0 to speed up tests
-    if "pytest" in sys.modules:
-        jitter = 0
-
-    return jitter
-
-
 class OAuth2Client:
     def __init__(self, credentials: Credentials, configuration=None):
         self._credentials = credentials

test/oauth2_test.py

@@ -11,7 +11,8 @@
 from openfga_sdk.constants import TOKEN_EXPIRY_THRESHOLD_BUFFER_IN_SEC, USER_AGENT
 from openfga_sdk.credentials import CredentialConfiguration, Credentials
 from openfga_sdk.exceptions import AuthenticationError
-from openfga_sdk.oauth2 import OAuth2Client, _TokenState
+from openfga_sdk.oauth2 import OAuth2Client
+from openfga_sdk.oauth2_common import _TokenState
 
 
 # Helper function to construct mock response

test/sync/oauth2_test.py

@@ -11,8 +11,9 @@
 from openfga_sdk.constants import TOKEN_EXPIRY_THRESHOLD_BUFFER_IN_SEC, USER_AGENT
 from openfga_sdk.credentials import CredentialConfiguration, Credentials
 from openfga_sdk.exceptions import AuthenticationError
+from openfga_sdk.oauth2_common import _TokenState
 from openfga_sdk.sync import rest
-from openfga_sdk.sync.oauth2 import OAuth2Client, _TokenState
+from openfga_sdk.sync.oauth2 import OAuth2Client
 
 
 # Helper function to construct mock response