Refactor access_token and a bit of cleanup

Author: Bre77

tesla_fleet_api/tesla/fleet.py

@@ -1,7 +1,7 @@
 """Tesla Fleet API for Python."""
 
 from json import dumps
-from typing import Any, Awaitable, Callable, Literal, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Literal
 
 import aiohttp
 
@@ -28,10 +28,9 @@
 class TeslaFleetApi(Tesla):
     """Class describing the Tesla Fleet API."""
 
-    access_token: str | Callable[[], Awaitable[str | None]] | None
+    _access_token: str | Callable[[], Awaitable[str | None]] | None
     server: str | None = None
     session: aiohttp.ClientSession
-    headers: dict[str, str]
     charging: "Charging"
     energySites: "EnergySites"
     user: "User"
@@ -40,8 +39,8 @@ class TeslaFleetApi(Tesla):
 
     def __init__(
         self,
-        session: aiohttp.ClientSession,
-        access_token: str | Callable[[], Awaitable[str | None]] | None,
+        session: aiohttp.ClientSession | None,
+        access_token: str | Callable[[], Awaitable[str | None]] | None = None,
         region: Literal["na", "eu", "cn"] | None = None,
         server: str | None = None,
         charging_scope: bool = True,
@@ -52,9 +51,8 @@ def __init__(
     ):
         """Initialize the Tesla Fleet API."""
 
-        self.session = session
-        self.access_token = access_token
-        self.headers = {}
+        self.session = session or aiohttp.ClientSession()
+        self._access_token = access_token
 
         if server is not None:
             self.server = server
@@ -92,12 +90,12 @@ async def find_server(self) -> str:
                 continue
         raise LibraryError("Could not find a valid Tesla API server.")
 
-    async def _access_token(self) -> str:
+    async def access_token(self) -> str:
         """Get the access token for the Tesla Fleet API."""
-        if callable(self.access_token):
-            token = await self.access_token()
+        if callable(self._access_token):
+            token = await self._access_token()
         else:
-            token = self.access_token
+            token = self._access_token
 
         if token is None:
             raise MissingToken
@@ -118,7 +116,7 @@ async def _request(
         if method == Method.GET:
             json = None
 
-        access_token = await self._access_token()
+        access_token = await self.access_token()
 
         headers = {
             "Authorization": f"Bearer {access_token}",
@@ -213,7 +211,7 @@ async def partner_login(
             if resp.ok:
                 token_data = await resp.json()
                 # Set the access token for subsequent API calls
-                self.access_token = token_data["access_token"]
+                self._access_token = token_data["access_token"]
                 return token_data
             else:
                 error_data = await resp.json()

tesla_fleet_api/tesla/oauth.py

@@ -1,9 +1,10 @@
+import time
 from typing import Any, Literal
+
 import aiohttp
-import time
 
+from tesla_fleet_api.const import SERVERS, Method, Scope
 from tesla_fleet_api.tesla import TeslaFleetApi
-from tesla_fleet_api.const import Scope, SERVERS, Method
 
 
 class TeslaFleetOAuth(TeslaFleetApi):
@@ -13,6 +14,7 @@ class TeslaFleetOAuth(TeslaFleetApi):
     refresh_token: str | None
     redirect_uri: str | None
     _client_secret: str | None
+    _access_token: str | None
 
     def __init__(
         self,
@@ -28,7 +30,7 @@ def __init__(
         self.client_id = client_id
         self._client_secret = client_secret
         self.redirect_uri = redirect_uri
-        self.access_token = access_token
+        self._access_token = access_token
         self.refresh_token = refresh_token
         self.expires = expires
 
@@ -75,14 +77,22 @@ async def get_refresh_token(self, code: str) -> None:
             if resp.ok:
                 data = await resp.json()
                 self.refresh_token = data.get("refresh_token")
-                self.access_token = data["access_token"]
+                self._access_token = data["access_token"]
                 self.expires = int(time.time()) + data["expires_in"]
                 region = code.split("_")[0].lower()
                 self.server = SERVERS.get(region)
 
+    async def access_token(self) -> str:
+        """Get the access token."""
+        if self.expires < time.time():
+            await self.refresh_access_token()
+        if self._access_token:
+            return self._access_token
+        raise ValueError("Access token is missing")
+
     async def check_access_token(self) -> dict[str, Any] | None:
         """Get the access token."""
-        if self.access_token and self.expires > time.time():
+        if self._access_token and self.expires > time.time():
             return None
         return await self.refresh_access_token()
 
@@ -100,7 +110,7 @@ async def refresh_access_token(self) -> dict[str, Any]:
         ) as resp:
             data = await resp.json()
             if resp.ok:
-                self.access_token = data["access_token"]
+                self._access_token = data["access_token"]
                 self.refresh_token = data["refresh_token"]
                 self.expires = int(time.time()) + data["expires_in"]
                 return {"refresh_token": self.refresh_token, "expires": self.expires}

tesla_fleet_api/teslemetry/teslemetry.py

@@ -21,7 +21,7 @@ def __init__(
         """Initialize the Teslemetry API."""
 
         self.session = session
-        self.access_token = access_token
+        self._access_token = access_token
         self.server = server
 
         self.charging = self.Charging(self)
@@ -110,7 +110,7 @@ async def migrate_to_oauth(
         self, client_id: str = "homeassistant", name: str | None = None
     ) -> dict[str, Any]:
         """Migrate from access token to OAuth."""
-        access_token = await self._access_token()
+        access_token = await self.access_token()
         migrate_data = {
             "grant_type": "migrate",
             "client_id": client_id,

tesla_fleet_api/tessie/tessie.py

@@ -1,13 +1,14 @@
-import aiohttp
 from typing import Any
 
-from tesla_fleet_api.tesla import TeslaFleetApi
+import aiohttp
+
 from tesla_fleet_api.const import Method
+from tesla_fleet_api.tesla import TeslaFleetApi
 from tesla_fleet_api.tessie.vehicles import TessieVehicles
 
-class Tessie(TeslaFleetApi):
 
-    server="https://api.tessie.com"
+class Tessie(TeslaFleetApi):
+    server = "https://api.tessie.com"
     Vehicles = TessieVehicles
 
     def __init__(
@@ -18,7 +19,7 @@ def __init__(
         """Initialize the Tessie API."""
 
         self.session = session
-        self.access_token = access_token
+        self._access_token = access_token
 
         self.charging = self.Charging(self)
         self.energySites = self.EnergySites(self)