Refactored code

src/blueapi/cli/cli.py

@@ -17,9 +17,13 @@
 from blueapi.client.client import BlueapiClient
 from blueapi.client.event_bus import AnyEvent, BlueskyStreamingError, EventBusClient
 from blueapi.client.rest import BlueskyRemoteControlError
-from blueapi.config import ApplicationConfig, CLIAuthConfig, ConfigLoader, OauthConfig
+from blueapi.config import (
+    ApplicationConfig,
+    CLIClientConfig,
+    ConfigLoader,
+)
 from blueapi.core import DataEvent
-from blueapi.service.authentication import TokenManager
+from blueapi.service.authentication import CLITokenManager, SessionManager
 from blueapi.service.main import start
 from blueapi.service.openapi import (
     DOCS_SCHEMA_LOCATION,
@@ -335,11 +339,13 @@ def scratch(obj: dict) -> None:
 @click.pass_obj
 def login(obj: dict) -> None:
     config: ApplicationConfig = obj["config"]
-    if config.cliAuth and config.oauth:
-        cliAuthConfig: CLIAuthConfig = config.cliAuth
-        oauthConfig: OauthConfig = config.oauth
+    if isinstance(config.oauth_client, CLIClientConfig) and config.oauth_server:
         print("Logging in")
-        auth: TokenManager = TokenManager(oauth=oauthConfig, cliAuth=cliAuthConfig)
+        auth: SessionManager = SessionManager(
+            server_config=config.oauth_server,
+            client_config=config.oauth_client,
+            token_manager=CLITokenManager(Path(config.oauth_client.token_file_path)),
+        )
         auth.start_device_flow()
     else:
         print("Please provide configuration to login!")
@@ -349,10 +355,12 @@ def login(obj: dict) -> None:
 @click.pass_obj
 def logout(obj: dict) -> None:
     config: ApplicationConfig = obj["config"]
-    if config.cliAuth and config.oauth:
-        oauthConfig: OauthConfig = config.oauth
-        cliAuthConfig: CLIAuthConfig = config.cliAuth
-        auth: TokenManager = TokenManager(cliAuth=cliAuthConfig, oauth=oauthConfig)
+    if isinstance(config.oauth_client, CLIClientConfig) and config.oauth_server:
+        auth: SessionManager = SessionManager(
+            server_config=config.oauth_server,
+            client_config=config.oauth_client,
+            token_manager=CLITokenManager(Path(config.oauth_client.token_file_path)),
+        )
         auth.logout()
         print("Logged out")
     else:

src/blueapi/client/client.py

@@ -6,6 +6,7 @@
 
 from blueapi.config import ApplicationConfig
 from blueapi.core.bluesky_types import DataEvent
+from blueapi.service.authentication import SessionManager
 from blueapi.service.model import (
     DeviceModel,
     DeviceResponse,
@@ -40,7 +41,8 @@ def __init__(
     @classmethod
     def from_config(cls, config: ApplicationConfig) -> "BlueapiClient":
         rest: BlueapiRestClient = BlueapiRestClient(
-            config.api, config.oauth, config.cliAuth
+            config.api,
+            SessionManager.from_config(config.oauth_server, config.oauth_client),
         )
         if config.stomp is not None:
             template = StompClient.for_broker(

src/blueapi/client/rest.py

@@ -5,8 +5,8 @@
 import requests
 from pydantic import TypeAdapter
 
-from blueapi.config import CLIAuthConfig, OauthConfig, RestConfig
-from blueapi.service.authentication import TokenManager
+from blueapi.config import RestConfig
+from blueapi.service.authentication import SessionManager
 from blueapi.service.model import (
     DeviceModel,
     DeviceResponse,
@@ -43,13 +43,10 @@ class BlueapiRestClient:
     def __init__(
         self,
         config: RestConfig | None = None,
-        authConfig: OauthConfig | None = None,
-        cliAuthConfig: CLIAuthConfig | None = None,
+        session_manager: SessionManager | None = None,
     ) -> None:
         self._config = config or RestConfig()
-        self._tokenHandler: TokenManager | None = None
-        if authConfig and cliAuthConfig:
-            self._tokenHandler = TokenManager(authConfig, cliAuthConfig)
+        self._session_manager: SessionManager | None = session_manager
 
     def get_plans(self) -> PlanResponse:
         return self._request_and_deserialize("/plans", PlanResponse)
@@ -140,19 +137,14 @@ def _request_and_deserialize(
         headers: dict[str, str] = {
             "content-type": "application/json; charset=UTF-8",
         }
-        if (
-            self._tokenHandler
-            and self._tokenHandler.token
-            and self._tokenHandler.token["access_token"]
-        ):
+        if self._session_manager and (token := self._session_manager.get_token()):
             try:
-                auth_token: str = self._tokenHandler.token["access_token"]
-                self._tokenHandler.authenticator.verify_token(auth_token)
-                headers["Authorization"] = f"Bearer {auth_token}"
+                self._session_manager.authenticator.verify_token(token["access_token"])
+                headers["Authorization"] = f"Bearer {token['access_token']}"
             except jwt.ExpiredSignatureError:
-                if self._tokenHandler.refresh_auth_token():
-                    access_token: str = self._tokenHandler.token["access_token"]
-                    headers["Authorization"] = f"Bearer {access_token}"
+                if token := self._session_manager.refresh_auth_token():
+                    if token := self._session_manager.get_token():
+                        headers["Authorization"] = f"Bearer {token['access_token']}"
             except Exception:
                 pass
         if data:

src/blueapi/config.py

@@ -83,7 +83,7 @@ class ScratchConfig(BlueapiBaseModel):
     repositories: list[ScratchRepository] = Field(default_factory=list)
 
 
-class OauthConfig(BlueapiBaseModel):
+class OAuthServerConfig(BlueapiBaseModel):
     oidc_config_url: str = Field(
         description="URL to fetch OIDC config from the provider"
     )
@@ -125,13 +125,13 @@ def model_post_init(self, __context: Any) -> None:
             raise ValueError("OIDC config is missing required fields")
 
 
-class BaseAuthConfig(BlueapiBaseModel):
+class OAuthClientConfig(BlueapiBaseModel):
     client_id: str = Field(description="Client ID")
     client_audience: str = Field(description="Client Audience")
 
 
-class CLIAuthConfig(BaseAuthConfig):
-    token_file_path: str = "~/token"
+class CLIClientConfig(OAuthClientConfig):
+    token_file_path: Path | None = Path("~/token")
 
 
 class ApplicationConfig(BlueapiBaseModel):
@@ -145,9 +145,8 @@ class ApplicationConfig(BlueapiBaseModel):
     logging: LoggingConfig = Field(default_factory=LoggingConfig)
     api: RestConfig = Field(default_factory=RestConfig)
     scratch: ScratchConfig | None = None
-    oauth: OauthConfig | None = None
-    cliAuth: CLIAuthConfig | None = None
-    swaggerAuth: BaseAuthConfig | None = None
+    oauth_server: OAuthServerConfig | None = None
+    oauth_client: CLIClientConfig | None = None
 
     def __eq__(self, other: object) -> bool:
         if isinstance(other, ApplicationConfig):

src/blueapi/service/__init__.py

@@ -1,4 +1,4 @@
-from .authentication import Authenticator, TokenManager
+from .authentication import Authenticator, SessionManager
 from .model import DeviceModel, PlanModel
 
-__all__ = ["PlanModel", "DeviceModel", "Authenticator", "TokenManager"]
+__all__ = ["PlanModel", "DeviceModel", "Authenticator", "SessionManager"]