From 3f9a08d295e7f61152bf15a4d3d0eb5aaebd3e44 Mon Sep 17 00:00:00 2001 From: Zoheb Shaikh <26975142+ZohebShaikh@users.noreply.github.com> Date: Thu, 10 Sep 2026 18:31:07 +0100 Subject: [PATCH] refactor: make client almost client again --- Dockerfile | 2 +- pyproject.toml | 31 +- src/blueapi/cli/cli.py | 2 +- src/blueapi/cli/scratch.py | 91 ------ src/blueapi/client/client.py | 2 +- src/blueapi/client/rest.py | 6 +- src/blueapi/client/session.py | 247 ++++++++++++++++ src/blueapi/core/context.py | 19 +- src/blueapi/log.py | 5 +- src/blueapi/service/__init__.py | 3 +- src/blueapi/service/authentication.py | 242 +--------------- src/blueapi/service/environment.py | 103 +++++++ src/blueapi/service/interface.py | 2 +- src/blueapi/{utils => service}/numtracker.py | 0 src/blueapi/utils/__init__.py | 2 - tests/system_tests/test_blueapi_system.py | 4 +- tests/unit_tests/cli/test_scratch.py | 267 +---------------- tests/unit_tests/client/test_rest.py | 2 +- tests/unit_tests/client/test_session.py | 120 ++++++++ .../unit_tests/service/test_authentication.py | 192 +++--------- tests/unit_tests/service/test_environment.py | 274 ++++++++++++++++++ tests/unit_tests/service/test_interface.py | 8 +- .../{utils => service}/test_numtracker.py | 2 +- tests/unit_tests/test_log.py | 4 +- uv.lock | 48 +-- 25 files changed, 872 insertions(+), 806 deletions(-) create mode 100644 src/blueapi/client/session.py create mode 100644 src/blueapi/service/environment.py rename src/blueapi/{utils => service}/numtracker.py (100%) create mode 100644 tests/unit_tests/client/test_session.py create mode 100644 tests/unit_tests/service/test_environment.py rename tests/unit_tests/{utils => service}/test_numtracker.py (99%) diff --git a/Dockerfile b/Dockerfile index 6311708419..09c9946593 100644 --- a/Dockerfile +++ b/Dockerfile @@ -29,7 +29,7 @@ ENV UV_PYTHON_INSTALL_DIR=/python # Sync the project without its dev dependencies RUN --mount=type=cache,target=/root/.cache/uv \ - uv sync --locked --no-editable --no-dev --managed-python + uv sync --locked --no-editable --no-dev --extra server --managed-python RUN uv pip install debugpy diff --git a/pyproject.toml b/pyproject.toml index e78835476e..3cd6631009 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,29 +14,25 @@ classifiers = [ ] description = "Lightweight bluesky-as-a-service wrapper application. Also usable as a library." dependencies = [ - "tiled[client]>=0.2.4", + # Base install is the `blueapi` CLI/client (and usable as a library on its own). + # Add blueapi[server] to also run `blueapi serve`. "bluesky[plotting]>=1.14.0", # plotting includes matplotlib, required for BestEffortCallback in run plans "ophyd-async>=0.13.5", - "aioca", "pydantic>=2.0", "pydantic-settings", "stomp-py", "PyYAML>=6.0.2", - "click>=8.2.0", - "fastapi>=0.112.0", - "uvicorn>=0.52.1", + "click>=8.2.0", # need for blueapi serve and blueapi setup-scratch "requests", - "GitPython>=3.1.58", #security base min https://github.com/advisories/GHSA-4gmw-gg2m-w46p + "GitPython>=3.1.58", #security base min https://github.com/advisories/GHSA-4gmw-gg2m-w46p - used by the server's /python_environment endpoint and the CLI's scratch setup "event-model==1.24.0", # https://github.com/DiamondLightSource/blueapi/issues/684 "bluesky-stomp>=0.1.6", "opentelemetry-distro>=0.48b0", - "opentelemetry-instrumentation-fastapi>=0.48b0", "observability-utils>=0.1.4", "pyjwt[crypto]", - "tomlkit", - "graypy>=2.1.0", - "httpx>=0.28.1", - "aiohttp>=3.13.5", + "websockets", + "tqdm", + "packaging", ] dynamic = ["version"] license.file = "LICENSE" @@ -45,9 +41,22 @@ requires-python = ">=3.11" [project.optional-dependencies] demo = ["dls-dodal>=1.69.0", "ophyd-async[sim]"] +# Only needed to run `blueapi serve` (the FastAPI worker service). +server = [ + "fastapi>=0.112.0", + "uvicorn>=0.52.1", + "aiohttp>=3.13.5", + "opentelemetry-instrumentation-fastapi>=0.48b0", + "graypy>=2.1.0", + "aioca", + "tiled[client]>=0.2.4", + "httpx>=0.28.1", # used by service/numtracker.py and service/authentication.py's TiledAuth + "tomlkit", +] [dependency-groups] dev = [ + "blueapi[server]", "ophyd_async[sim]", "copier", "dls-dodal>=1.69.0", diff --git a/src/blueapi/cli/cli.py b/src/blueapi/cli/cli.py index c3f9bde13f..394514eaa6 100644 --- a/src/blueapi/cli/cli.py +++ b/src/blueapi/cli/cli.py @@ -31,13 +31,13 @@ UnauthorisedAccessError, UnknownPlanError, ) +from blueapi.client.session import SessionCacheManager, SessionManager from blueapi.config import ( ApplicationConfig, ConfigLoader, ) from blueapi.core import OTLP_EXPORT_ENABLED, DataEvent from blueapi.log import set_up_logging -from blueapi.service.authentication import SessionCacheManager, SessionManager from blueapi.service.model import DeviceResponse, PlanResponse, SourceInfo, TaskRequest from blueapi.worker import ProgressEvent, WorkerEvent from blueapi.worker.event import TaskError, TaskResult diff --git a/src/blueapi/cli/scratch.py b/src/blueapi/cli/scratch.py index 6fa01d83e2..fd1367456d 100644 --- a/src/blueapi/cli/scratch.py +++ b/src/blueapi/cli/scratch.py @@ -1,4 +1,3 @@ -import importlib.metadata import logging import os import stat @@ -8,10 +7,8 @@ from subprocess import Popen from git import Repo -from tomlkit import parse from blueapi.config import FORBIDDEN_OWN_REMOTE_URL, ScratchConfig -from blueapi.service.model import PackageInfo, PythonEnvironmentResponse, SourceInfo from blueapi.utils import get_owner_gid, is_sgid_set _DEFAULT_INSTALL_TIMEOUT: float = 300.0 @@ -178,91 +175,3 @@ def _validate_directory(path: Path) -> None: raise KeyError(f"{path}: No such file or directory") elif path.is_file(): raise KeyError(f"{path}: Is a file, not a directory") - - -def _get_project_name_from_pyproject(path: Path) -> str: - pyproject_path = path / "pyproject.toml" - if pyproject_path.exists(): - with pyproject_path.open("r", encoding="utf-8") as file: - toml_data = parse(file.read()) - return toml_data.get("project", {}).get("name", "") - return "" - - -def _fetch_installed_packages_details() -> list[PackageInfo]: - installed_packages = importlib.metadata.distributions() - return [ - PackageInfo( - name=dist.metadata["Name"], - version=dist.version, - location=str(dist.locate_file("")), - is_dirty=False, - ) - for dist in installed_packages - ] - - -def get_python_environment( - config: ScratchConfig | None, - name: str | None = None, - source: SourceInfo | None = None, -) -> PythonEnvironmentResponse: - """ - Get the Python environment. This includes all installed packages and - the scratch packages. - """ - scratch_packages = {} - packages = [] - - if config is None: - python_env_response = PythonEnvironmentResponse(scratch_enabled=False) - else: - python_env_response = PythonEnvironmentResponse(scratch_enabled=True) - _validate_directory(config.root) - for repo in config.repositories: - local_directory = config.root / repo.name - repo = Repo(local_directory) - try: - branch = repo.active_branch.name - except TypeError: - branch = repo.head.commit.hexsha - - is_dirty = repo.is_dirty() - - version = ( - f"{repo.remotes[0].url} @{branch}" - if repo.remotes - else f"UNKNOWN REMOTE @{branch}" - ) - package_name = _get_project_name_from_pyproject(local_directory) - package_location = "" - - packages.append( - PackageInfo( - name=package_name, - version=version, - location=package_location, - source=SourceInfo.SCRATCH, - is_dirty=is_dirty, - ) - ) - scratch_packages = {p.name: p for p in packages} - - for pkg in _fetch_installed_packages_details(): - if pkg.name not in scratch_packages: - packages.append(pkg) - else: - scratch_packages[pkg.name].location += f"{pkg.location} &&" - - python_env_response.installed_packages = sorted( - packages, key=lambda pkg: pkg.name.lower() - ) - if name: - python_env_response.installed_packages = [ - p for p in python_env_response.installed_packages if p.name == name - ] - if source: - python_env_response.installed_packages = [ - p for p in python_env_response.installed_packages if p.source == source - ] - return python_env_response diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index 66a79d8dd6..e4860705eb 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -15,13 +15,13 @@ start_as_current_span, ) +from blueapi.client.session import SessionCacheManager, SessionManager from blueapi.config import ( ApplicationConfig, ConfigLoader, MissingStompConfigurationError, ) from blueapi.core.bluesky_types import DataEvent -from blueapi.service.authentication import SessionCacheManager, SessionManager from blueapi.service.model import ( DeviceModel, DeviceResponse, diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py index 6e03040c47..1389e1801f 100644 --- a/src/blueapi/client/rest.py +++ b/src/blueapi/client/rest.py @@ -1,10 +1,10 @@ import json import logging from collections.abc import Callable, Iterable, Mapping +from http import HTTPStatus from typing import Any, Literal, TypeVar import requests -from fastapi import status from observability_utils.tracing import ( get_context_propagator, get_tracer, @@ -17,9 +17,9 @@ from blueapi import __version__ from blueapi.client import client +from blueapi.client.session import JWTAuth, SessionManager from blueapi.config import RestConfig from blueapi.core.bluesky_types import DataEvent -from blueapi.service.authentication import JWTAuth, SessionManager from blueapi.service.model import ( DeviceModel, DeviceResponse, @@ -338,7 +338,7 @@ def _request_and_deserialize( exception = get_exception(response) if exception is not None: raise exception - if response.status_code == status.HTTP_204_NO_CONTENT: + if response.status_code == HTTPStatus.NO_CONTENT: raise NoContentError(target_type) if (server_version := response.headers.get("x-blueapi-version")) is not None: from packaging.version import Version diff --git a/src/blueapi/client/session.py b/src/blueapi/client/session.py new file mode 100644 index 0000000000..07ca6a7ae5 --- /dev/null +++ b/src/blueapi/client/session.py @@ -0,0 +1,247 @@ +from __future__ import annotations + +import base64 +import os +import time +import webbrowser +from abc import ABC, abstractmethod +from functools import cached_property +from http import HTTPStatus +from pathlib import Path +from typing import Any, cast + +import jwt +import requests +from pydantic import TypeAdapter +from requests.auth import AuthBase + +from blueapi.config import OIDCConfig +from blueapi.service.model import Cache + +DEFAULT_CACHE_DIR = "~/.cache/" +SCOPES = "openid offline_access" + + +class CacheManager(ABC): + @abstractmethod + def can_access_cache(self) -> bool: ... + @abstractmethod + def save_cache(self, cache: Cache) -> None: ... + @abstractmethod + def load_cache(self) -> Cache: ... + @abstractmethod + def delete_cache(self) -> None: ... + + +class SessionCacheManager(CacheManager): + def __init__(self, token_path: Path | None) -> None: + self._token_path: Path = ( + token_path if token_path else self._default_token_cache_path() + ) + + @cached_property + def _file_path(self) -> str: + return os.path.expanduser(self._token_path) + + def save_cache(self, cache: Cache) -> None: + self.delete_cache() + self._create_parent_folder_if_necessary() + with open(self._file_path, "xb") as token_file: + token_file.write(base64.b64encode(cache.model_dump_json().encode("utf-8"))) + os.chmod(self._file_path, 0o600) + + def load_cache(self) -> Cache: + self._create_parent_folder_if_necessary() + with open(self._file_path, "rb") as cache_file: + return TypeAdapter(Cache).validate_json( + base64.b64decode(cache_file.read()).decode("utf-8") + ) + + def delete_cache(self) -> None: + Path(self._file_path).unlink(missing_ok=True) + + @staticmethod + def _default_token_cache_path() -> Path: + """ + Return the default cache file path. + """ + cache_path = os.environ.get("XDG_CACHE_HOME", DEFAULT_CACHE_DIR) + return Path(cache_path).expanduser() / "blueapi_cache" + + def can_access_cache(self) -> bool: + assert self._token_path + try: + self._create_parent_folder_if_necessary() + self._token_path.write_text("") + except IsADirectoryError: + print("Invalid path: a directory path was provided instead of a file path") + return False + except PermissionError: + print(f"Permission denied: Cannot write to {self._token_path.absolute()}") + return False + return True + + def _create_parent_folder_if_necessary(self): + self._token_path.parent.mkdir(parents=True, exist_ok=True) + + +class SessionManager: + def __init__(self, server_config: OIDCConfig, cache_manager: CacheManager) -> None: + self._server_config = server_config + self._cache_manager: CacheManager = cache_manager + + @classmethod + def from_cache(cls, auth_token_path: Path | None) -> SessionManager: + cache_manager = SessionCacheManager(auth_token_path) + cache = cache_manager.load_cache() + return SessionManager( + server_config=cache.oidc_config, cache_manager=cache_manager + ) + + def delete_cache(self) -> None: + self._cache_manager.delete_cache() + + def get_valid_access_token(self) -> str: + """ + Retrieves a valid access token. + + Returns: + str: A valid access token if successful. + "": If the operation fails (no valid token could be fetched or refreshed) + """ + try: + cache = self._cache_manager.load_cache() + self.decode_jwt(cache.access_token) + return cache.access_token + except jwt.ExpiredSignatureError: + cache = self._cache_manager.load_cache() + return self._refresh_auth_token(cache.refresh_token) + except Exception: + self.delete_cache() + return "" + + @cached_property + def client(self): + return jwt.PyJWKClient(self._server_config.jwks_uri) + + def decode_jwt(self, json_web_token: str): + signing_key = self.client.get_signing_key_from_jwt(json_web_token) + return jwt.decode( + json_web_token, + signing_key.key, + algorithms=self._server_config.id_token_signing_alg_values_supported, + verify=True, + audience=self._server_config.client_audience, + issuer=self._server_config.issuer, + ) + + def logout(self) -> None: + cache = self._cache_manager.load_cache() + self.delete_cache() + try: + response = requests.get( + self._server_config.end_session_endpoint, + params={ + "id_token_hint": cache.id_token, + "client_id": self._server_config.client_id, + }, + ) + response.raise_for_status() + print("Logged out") + except Exception as e: + print( + "An unexpected error occurred while attempting " + f"to log out from the server.{e}" + ) + + def _refresh_auth_token(self, refresh_token: str) -> str: + response = requests.post( + self._server_config.token_endpoint, + data={ + "client_id": self._server_config.client_id, + "grant_type": "refresh_token", + "refresh_token": refresh_token, + }, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + ) + if response.status_code == HTTPStatus.OK: + token = response.json() + self._cache_manager.save_cache( + Cache( + oidc_config=self._server_config, + refresh_token=token["refresh_token"], + id_token=token["id_token"], + access_token=token["access_token"], + ), + ) + return token["access_token"] + else: + self.delete_cache() + return "" + + def poll_for_token( + self, device_code: str, polling_interval: float, expires_in: float + ) -> dict[str, Any]: + expiry_time: float = time.time() + expires_in + while time.time() < expiry_time: + response = requests.post( + self._server_config.token_endpoint, + data={ + "grant_type": "urn:ietf:params:oauth:grant-type:device_code", + "device_code": device_code, + "client_id": self._server_config.client_id, + }, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + ) + if response.status_code == HTTPStatus.OK: + return response.json() + time.sleep(polling_interval) + + raise TimeoutError("Polling timed out") + + def start_device_flow(self): + assert self._cache_manager.can_access_cache() + print("Logging in") + response: requests.Response = requests.post( + self._server_config.device_authorization_endpoint, + data={ + "client_id": self._server_config.client_id, + "scope": SCOPES, + }, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + ) + + response.raise_for_status() + + response_json: dict[str, Any] = response.json() + device_code = cast(str, response_json.get("device_code")) + interval = cast(float, response_json.get("interval")) + expires_in = cast(float, response_json.get("expires_in")) + webbrowser.open_new_tab(response_json["verification_uri_complete"]) + print( + f"Please login from this URL:- {response_json['verification_uri_complete']}" + ) + auth_token_json: dict[str, Any] = self.poll_for_token( + device_code, interval, expires_in + ) + self._cache_manager.save_cache( + Cache( + oidc_config=self._server_config, + refresh_token=auth_token_json["refresh_token"], + id_token=auth_token_json["id_token"], + access_token=auth_token_json["access_token"], + ) + ) + print("Logged in and cached new token") + + +class JWTAuth(AuthBase): + def __init__(self, session_manager: SessionManager | None): + self.token: str = ( + session_manager.get_valid_access_token() if session_manager else "" + ) + + def __call__(self, request): + if self.token: + request.headers["Authorization"] = f"Bearer {self.token}" + return request diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index 76ce74e647..bc5ece0629 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -5,7 +5,15 @@ from importlib import import_module, metadata from inspect import Parameter, isclass, signature from types import ModuleType, NoneType, UnionType -from typing import Any, TypeVar, Union, get_args, get_origin, get_type_hints +from typing import ( + TYPE_CHECKING, + Any, + TypeVar, + Union, + get_args, + get_origin, + get_type_hints, +) from bluesky.protocols import HasName from bluesky.run_engine import RunEngine @@ -32,13 +40,16 @@ from blueapi.core.protocols import DeviceManager from blueapi.utils import ( BlueapiPlanModelConfig, - NumtrackerClient, is_function_sourced_from_module, load_module_all, ) from blueapi.utils.invalid_config_error import InvalidConfigError from blueapi.utils.path_provider import StartDocumentPathProvider +if TYPE_CHECKING: + # Only needed by server config + from blueapi.service.numtracker import NumtrackerClient + from .bluesky_types import ( BLUESKY_PROTOCOLS, AsyncDevice, @@ -125,7 +136,7 @@ class BlueskyContext: default_factory=lambda: RunEngine(context_managers=[], call_returns_result=True) ) tiled_conf: TiledConfig | None = field(default=None, init=False, repr=False) - numtracker: NumtrackerClient | None = field(default=None, init=False, repr=False) + numtracker: "NumtrackerClient | None" = field(default=None, init=False, repr=False) path_provider: PathProvider | None = None plans: dict[str, Plan] = field(default_factory=dict) devices: dict[str, Device] = field(default_factory=dict) @@ -139,6 +150,8 @@ def __post_init__(self, configuration: ApplicationConfig | None): if (nt_conf := configuration.numtracker) is not None: if configuration.env.metadata is not None: + from blueapi.service.numtracker import NumtrackerClient + self.numtracker = NumtrackerClient(url=nt_conf.url) else: raise InvalidConfigError( diff --git a/src/blueapi/log.py b/src/blueapi/log.py index 1f28ad70bb..7c1406f060 100644 --- a/src/blueapi/log.py +++ b/src/blueapi/log.py @@ -5,7 +5,6 @@ from copy import copy import click -from graypy import GELFTCPHandler from blueapi.config import LoggingConfig @@ -107,13 +106,15 @@ def set_up_stream_handler( def set_up_graylog_handler( logger: logging.Logger, logging_config: LoggingConfig, filters: list[logging.Filter] -) -> GELFTCPHandler: +) -> logging.Handler: """Creates and configures GELFTCPHandler, then attaches to logger. Args: logger: Logger to attach handler to logging_config: LoggingConfig """ + from graypy import GELFTCPHandler + assert logging_config.graylog.url.host is not None, "Graylog URL missing host" assert logging_config.graylog.url.port is not None, "Graylog URL missing port" graylog_handler = GELFTCPHandler( diff --git a/src/blueapi/service/__init__.py b/src/blueapi/service/__init__.py index ad8ccb1f62..7c2fa404ca 100644 --- a/src/blueapi/service/__init__.py +++ b/src/blueapi/service/__init__.py @@ -1,4 +1,3 @@ -from .authentication import SessionManager from .model import DeviceModel, PlanModel -__all__ = ["PlanModel", "DeviceModel", "SessionManager"] +__all__ = ["PlanModel", "DeviceModel"] diff --git a/src/blueapi/service/authentication.py b/src/blueapi/service/authentication.py index 9177ad47fe..e5302774a3 100644 --- a/src/blueapi/service/authentication.py +++ b/src/blueapi/service/authentication.py @@ -1,257 +1,17 @@ from __future__ import annotations -import base64 -import os import threading -import time -import webbrowser -from abc import ABC, abstractmethod from collections.abc import Mapping -from functools import cached_property -from http import HTTPStatus -from pathlib import Path -from typing import Annotated, Any, cast +from typing import Annotated, Any import httpx import jwt -import requests from fastapi import Depends, HTTPException from fastapi.requests import HTTPConnection from fastapi.security.utils import get_authorization_scheme_param -from pydantic import TypeAdapter -from requests.auth import AuthBase from starlette.status import HTTP_401_UNAUTHORIZED from blueapi.config import OIDCConfig, ServiceAccount -from blueapi.service.model import Cache - -DEFAULT_CACHE_DIR = "~/.cache/" -SCOPES = "openid offline_access" - - -class CacheManager(ABC): - @abstractmethod - def can_access_cache(self) -> bool: ... - @abstractmethod - def save_cache(self, cache: Cache) -> None: ... - @abstractmethod - def load_cache(self) -> Cache: ... - @abstractmethod - def delete_cache(self) -> None: ... - - -class SessionCacheManager(CacheManager): - def __init__(self, token_path: Path | None) -> None: - self._token_path: Path = ( - token_path if token_path else self._default_token_cache_path() - ) - - @cached_property - def _file_path(self) -> str: - return os.path.expanduser(self._token_path) - - def save_cache(self, cache: Cache) -> None: - self.delete_cache() - self._create_parent_folder_if_necessary() - with open(self._file_path, "xb") as token_file: - token_file.write(base64.b64encode(cache.model_dump_json().encode("utf-8"))) - os.chmod(self._file_path, 0o600) - - def load_cache(self) -> Cache: - self._create_parent_folder_if_necessary() - with open(self._file_path, "rb") as cache_file: - return TypeAdapter(Cache).validate_json( - base64.b64decode(cache_file.read()).decode("utf-8") - ) - - def delete_cache(self) -> None: - Path(self._file_path).unlink(missing_ok=True) - - @staticmethod - def _default_token_cache_path() -> Path: - """ - Return the default cache file path. - """ - cache_path = os.environ.get("XDG_CACHE_HOME", DEFAULT_CACHE_DIR) - return Path(cache_path).expanduser() / "blueapi_cache" - - def can_access_cache(self) -> bool: - assert self._token_path - try: - self._create_parent_folder_if_necessary() - self._token_path.write_text("") - except IsADirectoryError: - print("Invalid path: a directory path was provided instead of a file path") - return False - except PermissionError: - print(f"Permission denied: Cannot write to {self._token_path.absolute()}") - return False - return True - - def _create_parent_folder_if_necessary(self): - self._token_path.parent.mkdir(parents=True, exist_ok=True) - - -class SessionManager: - def __init__(self, server_config: OIDCConfig, cache_manager: CacheManager) -> None: - self._server_config = server_config - self._cache_manager: CacheManager = cache_manager - - @classmethod - def from_cache(cls, auth_token_path: Path | None) -> SessionManager: - cache_manager = SessionCacheManager(auth_token_path) - cache = cache_manager.load_cache() - return SessionManager( - server_config=cache.oidc_config, cache_manager=cache_manager - ) - - def delete_cache(self) -> None: - self._cache_manager.delete_cache() - - def get_valid_access_token(self) -> str: - """ - Retrieves a valid access token. - - Returns: - str: A valid access token if successful. - "": If the operation fails (no valid token could be fetched or refreshed) - """ - try: - cache = self._cache_manager.load_cache() - self.decode_jwt(cache.access_token) - return cache.access_token - except jwt.ExpiredSignatureError: - cache = self._cache_manager.load_cache() - return self._refresh_auth_token(cache.refresh_token) - except Exception: - self.delete_cache() - return "" - - @cached_property - def client(self): - return jwt.PyJWKClient(self._server_config.jwks_uri) - - def decode_jwt(self, json_web_token: str): - signing_key = self.client.get_signing_key_from_jwt(json_web_token) - return jwt.decode( - json_web_token, - signing_key.key, - algorithms=self._server_config.id_token_signing_alg_values_supported, - verify=True, - audience=self._server_config.client_audience, - issuer=self._server_config.issuer, - ) - - def logout(self) -> None: - cache = self._cache_manager.load_cache() - self.delete_cache() - try: - response = requests.get( - self._server_config.end_session_endpoint, - params={ - "id_token_hint": cache.id_token, - "client_id": self._server_config.client_id, - }, - ) - response.raise_for_status() - print("Logged out") - except Exception as e: - print( - "An unexpected error occurred while attempting " - f"to log out from the server.{e}" - ) - - def _refresh_auth_token(self, refresh_token: str) -> str: - response = requests.post( - self._server_config.token_endpoint, - data={ - "client_id": self._server_config.client_id, - "grant_type": "refresh_token", - "refresh_token": refresh_token, - }, - headers={"Content-Type": "application/x-www-form-urlencoded"}, - ) - if response.status_code == HTTPStatus.OK: - token = response.json() - self._cache_manager.save_cache( - Cache( - oidc_config=self._server_config, - refresh_token=token["refresh_token"], - id_token=token["id_token"], - access_token=token["access_token"], - ), - ) - return token["access_token"] - else: - self.delete_cache() - return "" - - def poll_for_token( - self, device_code: str, polling_interval: float, expires_in: float - ) -> dict[str, Any]: - expiry_time: float = time.time() + expires_in - while time.time() < expiry_time: - response = requests.post( - self._server_config.token_endpoint, - data={ - "grant_type": "urn:ietf:params:oauth:grant-type:device_code", - "device_code": device_code, - "client_id": self._server_config.client_id, - }, - headers={"Content-Type": "application/x-www-form-urlencoded"}, - ) - if response.status_code == HTTPStatus.OK: - return response.json() - time.sleep(polling_interval) - - raise TimeoutError("Polling timed out") - - def start_device_flow(self): - assert self._cache_manager.can_access_cache() - print("Logging in") - response: requests.Response = requests.post( - self._server_config.device_authorization_endpoint, - data={ - "client_id": self._server_config.client_id, - "scope": SCOPES, - }, - headers={"Content-Type": "application/x-www-form-urlencoded"}, - ) - - response.raise_for_status() - - response_json: dict[str, Any] = response.json() - device_code = cast(str, response_json.get("device_code")) - interval = cast(float, response_json.get("interval")) - expires_in = cast(float, response_json.get("expires_in")) - webbrowser.open_new_tab(response_json["verification_uri_complete"]) - print( - f"Please login from this URL:- {response_json['verification_uri_complete']}" - ) - auth_token_json: dict[str, Any] = self.poll_for_token( - device_code, interval, expires_in - ) - self._cache_manager.save_cache( - Cache( - oidc_config=self._server_config, - refresh_token=auth_token_json["refresh_token"], - id_token=auth_token_json["id_token"], - access_token=auth_token_json["access_token"], - ) - ) - print("Logged in and cached new token") - - -class JWTAuth(AuthBase): - def __init__(self, session_manager: SessionManager | None): - self.token: str = ( - session_manager.get_valid_access_token() if session_manager else "" - ) - - def __call__(self, request): - if self.token: - request.headers["Authorization"] = f"Bearer {self.token}" - return request class TiledAuth(httpx.Auth): diff --git a/src/blueapi/service/environment.py b/src/blueapi/service/environment.py new file mode 100644 index 0000000000..af766a9a09 --- /dev/null +++ b/src/blueapi/service/environment.py @@ -0,0 +1,103 @@ +import importlib.metadata +from pathlib import Path + +from git import Repo +from tomlkit import parse + +from blueapi.config import ScratchConfig +from blueapi.service.model import PackageInfo, PythonEnvironmentResponse, SourceInfo + + +def _validate_directory(path: Path) -> None: + if not path.exists(): + raise KeyError(f"{path}: No such file or directory") + elif path.is_file(): + raise KeyError(f"{path}: Is a file, not a directory") + + +def _get_project_name_from_pyproject(path: Path) -> str: + pyproject_path = path / "pyproject.toml" + if pyproject_path.exists(): + with pyproject_path.open("r", encoding="utf-8") as file: + toml_data = parse(file.read()) + return toml_data.get("project", {}).get("name", "") + return "" + + +def _fetch_installed_packages_details() -> list[PackageInfo]: + installed_packages = importlib.metadata.distributions() + return [ + PackageInfo( + name=dist.metadata["Name"], + version=dist.version, + location=str(dist.locate_file("")), + is_dirty=False, + ) + for dist in installed_packages + ] + + +def get_python_environment( + config: ScratchConfig | None, + name: str | None = None, + source: SourceInfo | None = None, +) -> PythonEnvironmentResponse: + """ + Get the Python environment. This includes all installed packages and + the scratch packages. + """ + scratch_packages = {} + packages = [] + + if config is None: + python_env_response = PythonEnvironmentResponse(scratch_enabled=False) + else: + python_env_response = PythonEnvironmentResponse(scratch_enabled=True) + _validate_directory(config.root) + for repo in config.repositories: + local_directory = config.root / repo.name + repo = Repo(local_directory) + try: + branch = repo.active_branch.name + except TypeError: + branch = repo.head.commit.hexsha + + is_dirty = repo.is_dirty() + + version = ( + f"{repo.remotes[0].url} @{branch}" + if repo.remotes + else f"UNKNOWN REMOTE @{branch}" + ) + package_name = _get_project_name_from_pyproject(local_directory) + package_location = "" + + packages.append( + PackageInfo( + name=package_name, + version=version, + location=package_location, + source=SourceInfo.SCRATCH, + is_dirty=is_dirty, + ) + ) + scratch_packages = {p.name: p for p in packages} + + for pkg in _fetch_installed_packages_details(): + if pkg.name not in scratch_packages: + packages.append(pkg) + else: + scratch_packages[pkg.name].location += f"{pkg.location} &&" + + python_env_response.installed_packages = sorted( + packages, key=lambda pkg: pkg.name.lower() + ) + if name: + python_env_response.installed_packages = [ + p for p in python_env_response.installed_packages if p.name == name + ] + if source: + python_env_response.installed_packages = [ + p for p in python_env_response.installed_packages if p.source == source + ] + return python_env_response diff --git a/src/blueapi/service/interface.py b/src/blueapi/service/interface.py index 6c0c5befbd..45dce059ba 100644 --- a/src/blueapi/service/interface.py +++ b/src/blueapi/service/interface.py @@ -10,13 +10,13 @@ from bluesky_stomp.models import Broker, DestinationBase, MessageTopic from tiled.client import from_uri -from blueapi.cli.scratch import get_python_environment from blueapi.config import ApplicationConfig, OIDCConfig, ServiceAccount, StompConfig from blueapi.core.bluesky_types import DataEvent from blueapi.core.context import BlueskyContext from blueapi.core.event import EventStream from blueapi.log import set_up_logging from blueapi.service.authentication import TiledAuth +from blueapi.service.environment import get_python_environment from blueapi.service.model import ( DeviceModel, PlanModel, diff --git a/src/blueapi/utils/numtracker.py b/src/blueapi/service/numtracker.py similarity index 100% rename from src/blueapi/utils/numtracker.py rename to src/blueapi/service/numtracker.py diff --git a/src/blueapi/utils/__init__.py b/src/blueapi/utils/__init__.py index f722c5b42d..4a0271851d 100644 --- a/src/blueapi/utils/__init__.py +++ b/src/blueapi/utils/__init__.py @@ -8,7 +8,6 @@ from .file_permissions import get_owner_gid, is_sgid_set from .invalid_config_error import InvalidConfigError from .modules import is_function_sourced_from_module, load_module_all -from .numtracker import NumtrackerClient from .serialization import serialize from .thread_exception import handle_all_exceptions @@ -20,7 +19,6 @@ "BlueapiModelConfig", "BlueapiPlanModelConfig", "InvalidConfigError", - "NumtrackerClient", "report_successful_devices", "is_sgid_set", "get_owner_gid", diff --git a/tests/system_tests/test_blueapi_system.py b/tests/system_tests/test_blueapi_system.py index 8c4dc19cf9..7d56aa7f1b 100644 --- a/tests/system_tests/test_blueapi_system.py +++ b/tests/system_tests/test_blueapi_system.py @@ -156,7 +156,7 @@ def get_access_token(user: ValidUser) -> str: @pytest.fixture(scope="module") def client_without_auth() -> Generator[BlueapiClient]: with patch( - "blueapi.service.authentication.SessionManager.from_cache", + "blueapi.client.session.SessionManager.from_cache", return_value=None, ): yield BlueapiClient.from_config(config=ApplicationConfig()) @@ -166,7 +166,7 @@ def patch_session(user: ValidUser): mock_session_manager = MagicMock() mock_session_manager.get_valid_access_token.return_value = get_access_token(user) return patch( - "blueapi.service.authentication.SessionManager.from_cache", + "blueapi.client.session.SessionManager.from_cache", return_value=mock_session_manager, ) diff --git a/tests/unit_tests/cli/test_scratch.py b/tests/unit_tests/cli/test_scratch.py index 22f3ad3d8f..a7df496186 100644 --- a/tests/unit_tests/cli/test_scratch.py +++ b/tests/unit_tests/cli/test_scratch.py @@ -3,21 +3,13 @@ import uuid from collections.abc import Generator from pathlib import Path -from unittest.mock import ANY, MagicMock, Mock, PropertyMock, call, patch +from unittest.mock import ANY, MagicMock, Mock, call, patch import pytest from git import Repo -from blueapi.cli.scratch import ( - _fetch_installed_packages_details, - _get_project_name_from_pyproject, - ensure_repo, - get_python_environment, - scratch_install, - setup_scratch, -) +from blueapi.cli.scratch import ensure_repo, scratch_install, setup_scratch from blueapi.config import ScratchConfig, ScratchRepository -from blueapi.service.model import PackageInfo, SourceInfo from blueapi.utils import get_owner_gid @@ -383,258 +375,3 @@ def test_setup_scratch_continues_after_failure( RuntimeError, match="Failed to clone", check=lambda e: str(e.__cause__) == "bar" ): setup_scratch(config) - - -@pytest.fixture -def config(directory_path_with_sgid: Path) -> ScratchConfig: - return ScratchConfig( - root=directory_path_with_sgid, - repositories=[ - ScratchRepository( - name="foo", - remote_url="http://example.com/foo.git", - ), - ScratchRepository( - name="bar", - remote_url="http://example.com/bar.git", - ), - ], - ) - - -@patch("blueapi.cli.scratch.Repo") -@patch("blueapi.cli.scratch._fetch_installed_packages_details") -@patch("blueapi.cli.scratch._get_project_name_from_pyproject") -def test_get_python_env_returns_correct_packages( - mock_get_project_name: Mock, - mock_fetch_installed_packages: Mock, - mock_repo: Mock, - directory_path_with_sgid: Path, - config: ScratchConfig, -): - repo_path = directory_path_with_sgid / "foo" - repo_path.mkdir() - mock_repo_1 = Mock() - mock_repo_1.active_branch.name = "main" - mock_repo_1.is_dirty.return_value = False - mock_repo_1.remotes = [Mock(url="http://example.com/foo.git")] - - repo_path = directory_path_with_sgid / "bar" - repo_path.mkdir() - mock_repo_2 = Mock() - type(mock_repo_2.active_branch).name = PropertyMock(side_effect=TypeError) - mock_repo_2.head.commit.hexsha = "adsad23123" - mock_repo_2.is_dirty.return_value = True - mock_repo_2.remotes = [Mock(url="http://example.com/bar.git")] - - mock_repo.side_effect = [mock_repo_1, mock_repo_2] - - mock_get_project_name.side_effect = ["foo-package", "bar-package"] - mock_fetch_installed_packages.return_value = [ - PackageInfo( - name="package-01", - version="1.0.1", - location="/some/location", - is_dirty=False, - ) - ] - - response = get_python_environment(config) - - assert response.installed_packages == [ - PackageInfo( - name="bar-package", - version="http://example.com/bar.git @adsad23123", - location="", - is_dirty=True, - source=SourceInfo.SCRATCH, - ), - PackageInfo( - name="foo-package", - version="http://example.com/foo.git @main", - location="", - is_dirty=False, - source=SourceInfo.SCRATCH, - ), - PackageInfo( - name="package-01", - version="1.0.1", - location="/some/location", - is_dirty=False, - source=SourceInfo.PYPI, - ), - ] - - -@patch("blueapi.cli.scratch.Repo") -@patch("blueapi.cli.scratch._fetch_installed_packages_details") -@patch("blueapi.cli.scratch._get_project_name_from_pyproject") -def test_fetch_python_env_with_identical_packages( - mock_get_project_name: Mock, - mock_fetch_installed_packages: Mock, - mock_repo: Mock, - directory_path_with_sgid: Path, -): - repo_path = directory_path_with_sgid / "foo" - repo_path.mkdir() - mock_repo_instance = Mock() - mock_repo_instance.active_branch.name = "main" - mock_repo_instance.is_dirty.return_value = False - mock_repo_instance.remotes = [Mock(url="http://example.com/foo.git")] - - mock_repo.return_value = mock_repo_instance - - mock_get_project_name.return_value = "foo-package" - mock_fetch_installed_packages.return_value = [ - PackageInfo( - name="foo-package", - version="http://example.com/foo.git @main", - location="/some/location", - is_dirty=False, - source=SourceInfo.SCRATCH, - ) - ] - config = ScratchConfig( - root=directory_path_with_sgid, - repositories=[ - ScratchRepository( - name="foo", - remote_url="http://example.com/foo.git", - ), - ], - ) - response = get_python_environment(config) - - assert response.installed_packages == [ - PackageInfo( - name="foo-package", - version="http://example.com/foo.git @main", - location="/some/location &&", - is_dirty=False, - source=SourceInfo.SCRATCH, - ), - ] - - -@patch("blueapi.cli.scratch.importlib.metadata.distributions") -def test_fetch_installed_packages_details_returns_correct_packages(mock_distributions): - mock_distribution = Mock() - mock_distribution.metadata = {"Name": "example-package"} - mock_distribution.version = "1.0.0" - mock_distribution.locate_file.return_value = Path("/example/location") - mock_distributions.return_value = [mock_distribution] - - packages = _fetch_installed_packages_details() - - assert len(packages) == 1 - assert packages == [ - PackageInfo( - name="example-package", - version="1.0.0", - location="/example/location", - is_dirty=False, - ) - ] - - -@patch("blueapi.cli.scratch.Repo") -@patch("blueapi.cli.scratch._fetch_installed_packages_details") -@patch("blueapi.cli.scratch._get_project_name_from_pyproject") -def test_get_python_env_filters_by_name_and_source( - mock_get_project_name: Mock, - mock_fetch_installed_packages: Mock, - mock_repo: Mock, - directory_path_with_sgid: Path, -): - # Setup for scratch source filtering - repo_path = directory_path_with_sgid / "foo" - repo_path.mkdir() - mock_repo_instance = Mock() - mock_repo_instance.active_branch.name = "main" - mock_repo_instance.is_dirty.return_value = False - mock_repo_instance.remotes = [Mock(url="http://example.com/foo.git")] - mock_repo.return_value = mock_repo_instance - - mock_get_project_name.return_value = "foo-package" - mock_fetch_installed_packages.return_value = [ - PackageInfo( - name="bar-package", - version="1.0.0", - location="/some/location", - is_dirty=False, - source=SourceInfo.PYPI, - ) - ] - config = ScratchConfig( - root=directory_path_with_sgid, - repositories=[ - ScratchRepository( - name="foo", - remote_url="http://example.com/foo.git", - ), - ], - ) - # Test filtering by name - response_by_name = get_python_environment(config, name="foo-package") - assert response_by_name.installed_packages == [ - PackageInfo( - name="foo-package", - version="http://example.com/foo.git @main", - location="", - is_dirty=False, - source=SourceInfo.SCRATCH, - ) - ] - - # Test filtering by source - response_by_source = get_python_environment(config, source=SourceInfo.SCRATCH) - assert response_by_source.installed_packages == [ - PackageInfo( - name="foo-package", - version="http://example.com/foo.git @main", - location="", - is_dirty=False, - source=SourceInfo.SCRATCH, - ) - ] - - -@pytest.fixture -def pyproject_file(tmp_path: Path) -> Generator[Path]: - pyproject_path = tmp_path / "pyproject.toml" - with pyproject_path.open("w") as f: - f.write( - """ - [project] - name = "example-project" - """ - ) - yield pyproject_path - os.remove(pyproject_path) - - -def test_get_project_name_from_pyproject_returns_name(pyproject_file: Path): - project_name = _get_project_name_from_pyproject(pyproject_file.parent) - assert project_name == "example-project" - - -def test_get_project_name_from_pyproject_returns_empty_if_no_pyproject( - tmp_path: Path, -): - project_name = _get_project_name_from_pyproject(tmp_path) - assert project_name == "" - - -def test_get_project_name_from_pyproject_returns_empty_if_no_name_key( - tmp_path: Path, -): - pyproject_path = tmp_path / "pyproject.toml" - with pyproject_path.open("w") as f: - f.write( - """ - [project] - version = "1.0.0" - """ - ) - project_name = _get_project_name_from_pyproject(tmp_path) - assert project_name == "" diff --git a/tests/unit_tests/client/test_rest.py b/tests/unit_tests/client/test_rest.py index 6ecfbaa765..4b60f0fb05 100644 --- a/tests/unit_tests/client/test_rest.py +++ b/tests/unit_tests/client/test_rest.py @@ -28,8 +28,8 @@ _create_task_exceptions, _exception, ) +from blueapi.client.session import SessionCacheManager, SessionManager from blueapi.config import OIDCConfig -from blueapi.service.authentication import SessionCacheManager, SessionManager from blueapi.service.model import ( DeviceModel, EnvironmentResponse, diff --git a/tests/unit_tests/client/test_session.py b/tests/unit_tests/client/test_session.py new file mode 100644 index 0000000000..f54334fdc4 --- /dev/null +++ b/tests/unit_tests/client/test_session.py @@ -0,0 +1,120 @@ +import os +from pathlib import Path +from typing import Any +from unittest.mock import patch + +import pytest +import responses +from starlette.status import HTTP_403_FORBIDDEN + +from blueapi.client.session import SessionCacheManager, SessionManager +from blueapi.config import OIDCConfig + + +@pytest.fixture +def auth_token_path(tmp_path) -> Path: + return tmp_path / "blueapi_cache" + + +@pytest.fixture +def session_manager( + oidc_config: OIDCConfig, + auth_token_path, + mock_authn_server: responses.RequestsMock, +) -> SessionManager: + return SessionManager( + server_config=oidc_config, cache_manager=SessionCacheManager(auth_token_path) + ) + + +def test_logout( + session_manager: SessionManager, + oidc_config: OIDCConfig, + cached_valid_token: Path, + auth_token_path: Path, +): + assert os.path.exists(auth_token_path) + session_manager.logout() + assert not os.path.exists(auth_token_path) + + +def test_refresh_auth_token( + mock_authn_server: responses.RequestsMock, + session_manager: SessionManager, + cached_valid_refresh: Path, +): + token = session_manager.get_valid_access_token() + assert token == "new_token" + + +def test_get_empty_token_if_no_cache(session_manager: SessionManager): + token = session_manager.get_valid_access_token() + assert token == "" + + +def test_get_empty_token_if_refresh_fails( + mock_authn_server: responses.RequestsMock, + session_manager: SessionManager, + cached_expired_refresh: Path, +): + assert cached_expired_refresh.exists() + token = session_manager.get_valid_access_token() + assert token == "" + assert not cached_expired_refresh.exists() + + +def test_get_empty_token_if_invalid_cache( + mock_authn_server: responses.RequestsMock, + session_manager: SessionManager, + cache_with_invalid_audience: Path, +): + token = session_manager.get_valid_access_token() + assert token == "" + + +def test_get_empty_token_if_exception_in_decode( + mock_authn_server: responses.RequestsMock, + session_manager: SessionManager, + cached_expired_refresh: Path, +): + assert cached_expired_refresh.exists() + token = session_manager.get_valid_access_token() + assert token == "" + assert not cached_expired_refresh.exists() + + +def test_poll_for_token( + mock_authn_server: responses.RequestsMock, + session_manager: SessionManager, + valid_token: dict[str, Any], + device_code: str, +): + token = session_manager.poll_for_token(device_code, 1, 2) + assert token == valid_token + + +@patch("blueapi.client.session.time.sleep", return_value=None) +def test_poll_for_token_timeout( + mock_sleep, + oidc_well_known: dict[str, Any], + mock_authn_server: responses.RequestsMock, + session_manager: SessionManager, + device_code: str, +): + mock_authn_server.stop() + mock_authn_server.remove(responses.POST, oidc_well_known["token_endpoint"]) + mock_authn_server.post( + url=oidc_well_known["token_endpoint"], + json={"error": "authorization_pending"}, + status=HTTP_403_FORBIDDEN, + ) + with pytest.raises(TimeoutError), mock_authn_server: + session_manager.poll_for_token(device_code, 0.01, 0.01) + + +def test_session_cache_manager_returns_writable_file_path(tmp_path): + os.environ["XDG_CACHE_HOME"] = str(tmp_path) + cache = SessionCacheManager(token_path=None) + Path(cache._file_path).touch() + assert os.path.isfile(cache._file_path) + assert cache._file_path == f"{tmp_path}/blueapi_cache" diff --git a/tests/unit_tests/service/test_authentication.py b/tests/unit_tests/service/test_authentication.py index a763758122..c496164873 100644 --- a/tests/unit_tests/service/test_authentication.py +++ b/tests/unit_tests/service/test_authentication.py @@ -1,6 +1,3 @@ -import os -from pathlib import Path -from typing import Any from unittest.mock import Mock, patch import httpx @@ -10,13 +7,11 @@ import respx from fastapi import HTTPException from pydantic import SecretStr -from starlette.status import HTTP_200_OK, HTTP_403_FORBIDDEN +from starlette.status import HTTP_200_OK from blueapi.config import OIDCConfig, ServiceAccount from blueapi.service import authentication from blueapi.service.authentication import ( - SessionCacheManager, - SessionManager, TiledAuth, access_token, build_access_token_check, @@ -24,107 +19,6 @@ ) -@pytest.fixture -def auth_token_path(tmp_path) -> Path: - return tmp_path / "blueapi_cache" - - -@pytest.fixture -def session_manager( - oidc_config: OIDCConfig, - auth_token_path, - mock_authn_server: responses.RequestsMock, -) -> SessionManager: - return SessionManager( - server_config=oidc_config, cache_manager=SessionCacheManager(auth_token_path) - ) - - -def test_logout( - session_manager: SessionManager, - oidc_config: OIDCConfig, - cached_valid_token: Path, - auth_token_path: Path, -): - assert os.path.exists(auth_token_path) - session_manager.logout() - assert not os.path.exists(auth_token_path) - - -def test_refresh_auth_token( - mock_authn_server: responses.RequestsMock, - session_manager: SessionManager, - cached_valid_refresh: Path, -): - token = session_manager.get_valid_access_token() - assert token == "new_token" - - -def test_get_empty_token_if_no_cache(session_manager: SessionManager): - token = session_manager.get_valid_access_token() - assert token == "" - - -def test_get_empty_token_if_refresh_fails( - mock_authn_server: responses.RequestsMock, - session_manager: SessionManager, - cached_expired_refresh: Path, -): - assert cached_expired_refresh.exists() - token = session_manager.get_valid_access_token() - assert token == "" - assert not cached_expired_refresh.exists() - - -def test_get_empty_token_if_invalid_cache( - mock_authn_server: responses.RequestsMock, - session_manager: SessionManager, - cache_with_invalid_audience: Path, -): - token = session_manager.get_valid_access_token() - assert token == "" - - -def test_get_empty_token_if_exception_in_decode( - mock_authn_server: responses.RequestsMock, - session_manager: SessionManager, - cached_expired_refresh: Path, -): - assert cached_expired_refresh.exists() - token = session_manager.get_valid_access_token() - assert token == "" - assert not cached_expired_refresh.exists() - - -def test_poll_for_token( - mock_authn_server: responses.RequestsMock, - session_manager: SessionManager, - valid_token: dict[str, Any], - device_code: str, -): - token = session_manager.poll_for_token(device_code, 1, 2) - assert token == valid_token - - -@patch("blueapi.service.authentication.time.sleep", return_value=None) -def test_poll_for_token_timeout( - mock_sleep, - oidc_well_known: dict[str, Any], - mock_authn_server: responses.RequestsMock, - session_manager: SessionManager, - device_code: str, -): - mock_authn_server.stop() - mock_authn_server.remove(responses.POST, oidc_well_known["token_endpoint"]) - mock_authn_server.post( - url=oidc_well_known["token_endpoint"], - json={"error": "authorization_pending"}, - status=HTTP_403_FORBIDDEN, - ) - with pytest.raises(TimeoutError), mock_authn_server: - session_manager.poll_for_token(device_code, 0.01, 0.01) - - def test_server_raises_exception_for_invalid_token( oidc_config: OIDCConfig, mock_authn_server: responses.RequestsMock ): @@ -142,52 +36,6 @@ def test_processes_valid_token( inner(Mock(), token=valid_token_with_jwt["access_token"]) -def test_session_cache_manager_returns_writable_file_path(tmp_path): - os.environ["XDG_CACHE_HOME"] = str(tmp_path) - cache = SessionCacheManager(token_path=None) - Path(cache._file_path).touch() - assert os.path.isfile(cache._file_path) - assert cache._file_path == f"{tmp_path}/blueapi_cache" - - -def test_tiled_auth_raises_exception(): - with pytest.raises( - RuntimeError, match="Token URL is not set please check oidc config" - ): - auth = ServiceAccount() - TiledAuth(tiled_auth=auth) - - -@respx.mock -def test_tiled_auth_sync_auth_flow(): - client_id = "client" - client_secret = SecretStr("secret") - token_url = "http://keycloak.com/token" - access_token = "access_token" - - respx.post(token_url).mock( - return_value=httpx.Response( - status_code=HTTP_200_OK, json={"access_token": access_token} - ) - ) - - tiled_auth = TiledAuth( - tiled_auth=ServiceAccount( - client_id=client_id, - client_secret=client_secret, - token_url=token_url, - ) - ) - - request = Mock() - request.headers = {} - - flow = tiled_auth.sync_auth_flow(request) - result = next(flow) - - assert result.headers["Authorization"] == f"Bearer {access_token}" - - @pytest.mark.parametrize( "header,cookie,token", [ @@ -236,3 +84,41 @@ def test_build_access_token(mock_jwt: Mock): validate_fn(req, token=None) mock_jwt.decode.assert_not_called() + + +def test_tiled_auth_raises_exception(): + with pytest.raises( + RuntimeError, match="Token URL is not set please check oidc config" + ): + auth = ServiceAccount() + TiledAuth(tiled_auth=auth) + + +@respx.mock +def test_tiled_auth_sync_auth_flow(): + client_id = "client" + client_secret = SecretStr("secret") + token_url = "http://keycloak.com/token" + access_token = "access_token" + + respx.post(token_url).mock( + return_value=httpx.Response( + status_code=HTTP_200_OK, json={"access_token": access_token} + ) + ) + + tiled_auth = TiledAuth( + tiled_auth=ServiceAccount( + client_id=client_id, + client_secret=client_secret, + token_url=token_url, + ) + ) + + request = Mock() + request.headers = {} + + flow = tiled_auth.sync_auth_flow(request) + result = next(flow) + + assert result.headers["Authorization"] == f"Bearer {access_token}" diff --git a/tests/unit_tests/service/test_environment.py b/tests/unit_tests/service/test_environment.py new file mode 100644 index 0000000000..5ce15b283f --- /dev/null +++ b/tests/unit_tests/service/test_environment.py @@ -0,0 +1,274 @@ +import os +from collections.abc import Generator +from pathlib import Path +from unittest.mock import Mock, PropertyMock, patch + +import pytest + +from blueapi.config import ScratchConfig, ScratchRepository +from blueapi.service.environment import ( + _fetch_installed_packages_details, + _get_project_name_from_pyproject, + get_python_environment, +) +from blueapi.service.model import PackageInfo, SourceInfo + + +@pytest.fixture +def directory_path_with_sgid(tmp_path: Path) -> Path: + return tmp_path + + +@pytest.fixture +def config(directory_path_with_sgid: Path) -> ScratchConfig: + return ScratchConfig( + root=directory_path_with_sgid, + repositories=[ + ScratchRepository( + name="foo", + remote_url="http://example.com/foo.git", + ), + ScratchRepository( + name="bar", + remote_url="http://example.com/bar.git", + ), + ], + ) + + +@patch("blueapi.service.environment.Repo") +@patch("blueapi.service.environment._fetch_installed_packages_details") +@patch("blueapi.service.environment._get_project_name_from_pyproject") +def test_get_python_env_returns_correct_packages( + mock_get_project_name: Mock, + mock_fetch_installed_packages: Mock, + mock_repo: Mock, + directory_path_with_sgid: Path, + config: ScratchConfig, +): + repo_path = directory_path_with_sgid / "foo" + repo_path.mkdir() + mock_repo_1 = Mock() + mock_repo_1.active_branch.name = "main" + mock_repo_1.is_dirty.return_value = False + mock_repo_1.remotes = [Mock(url="http://example.com/foo.git")] + + repo_path = directory_path_with_sgid / "bar" + repo_path.mkdir() + mock_repo_2 = Mock() + type(mock_repo_2.active_branch).name = PropertyMock(side_effect=TypeError) + mock_repo_2.head.commit.hexsha = "adsad23123" + mock_repo_2.is_dirty.return_value = True + mock_repo_2.remotes = [Mock(url="http://example.com/bar.git")] + + mock_repo.side_effect = [mock_repo_1, mock_repo_2] + + mock_get_project_name.side_effect = ["foo-package", "bar-package"] + mock_fetch_installed_packages.return_value = [ + PackageInfo( + name="package-01", + version="1.0.1", + location="/some/location", + is_dirty=False, + ) + ] + + response = get_python_environment(config) + + assert response.installed_packages == [ + PackageInfo( + name="bar-package", + version="http://example.com/bar.git @adsad23123", + location="", + is_dirty=True, + source=SourceInfo.SCRATCH, + ), + PackageInfo( + name="foo-package", + version="http://example.com/foo.git @main", + location="", + is_dirty=False, + source=SourceInfo.SCRATCH, + ), + PackageInfo( + name="package-01", + version="1.0.1", + location="/some/location", + is_dirty=False, + source=SourceInfo.PYPI, + ), + ] + + +@patch("blueapi.service.environment.Repo") +@patch("blueapi.service.environment._fetch_installed_packages_details") +@patch("blueapi.service.environment._get_project_name_from_pyproject") +def test_fetch_python_env_with_identical_packages( + mock_get_project_name: Mock, + mock_fetch_installed_packages: Mock, + mock_repo: Mock, + directory_path_with_sgid: Path, +): + repo_path = directory_path_with_sgid / "foo" + repo_path.mkdir() + mock_repo_instance = Mock() + mock_repo_instance.active_branch.name = "main" + mock_repo_instance.is_dirty.return_value = False + mock_repo_instance.remotes = [Mock(url="http://example.com/foo.git")] + + mock_repo.return_value = mock_repo_instance + + mock_get_project_name.return_value = "foo-package" + mock_fetch_installed_packages.return_value = [ + PackageInfo( + name="foo-package", + version="http://example.com/foo.git @main", + location="/some/location", + is_dirty=False, + source=SourceInfo.SCRATCH, + ) + ] + config = ScratchConfig( + root=directory_path_with_sgid, + repositories=[ + ScratchRepository( + name="foo", + remote_url="http://example.com/foo.git", + ), + ], + ) + response = get_python_environment(config) + + assert response.installed_packages == [ + PackageInfo( + name="foo-package", + version="http://example.com/foo.git @main", + location="/some/location &&", + is_dirty=False, + source=SourceInfo.SCRATCH, + ), + ] + + +@patch("blueapi.service.environment.importlib.metadata.distributions") +def test_fetch_installed_packages_details_returns_correct_packages(mock_distributions): + mock_distribution = Mock() + mock_distribution.metadata = {"Name": "example-package"} + mock_distribution.version = "1.0.0" + mock_distribution.locate_file.return_value = Path("/example/location") + mock_distributions.return_value = [mock_distribution] + + packages = _fetch_installed_packages_details() + + assert len(packages) == 1 + assert packages == [ + PackageInfo( + name="example-package", + version="1.0.0", + location="/example/location", + is_dirty=False, + ) + ] + + +@patch("blueapi.service.environment.Repo") +@patch("blueapi.service.environment._fetch_installed_packages_details") +@patch("blueapi.service.environment._get_project_name_from_pyproject") +def test_get_python_env_filters_by_name_and_source( + mock_get_project_name: Mock, + mock_fetch_installed_packages: Mock, + mock_repo: Mock, + directory_path_with_sgid: Path, +): + # Setup for scratch source filtering + repo_path = directory_path_with_sgid / "foo" + repo_path.mkdir() + mock_repo_instance = Mock() + mock_repo_instance.active_branch.name = "main" + mock_repo_instance.is_dirty.return_value = False + mock_repo_instance.remotes = [Mock(url="http://example.com/foo.git")] + mock_repo.return_value = mock_repo_instance + + mock_get_project_name.return_value = "foo-package" + mock_fetch_installed_packages.return_value = [ + PackageInfo( + name="bar-package", + version="1.0.0", + location="/some/location", + is_dirty=False, + source=SourceInfo.PYPI, + ) + ] + config = ScratchConfig( + root=directory_path_with_sgid, + repositories=[ + ScratchRepository( + name="foo", + remote_url="http://example.com/foo.git", + ), + ], + ) + # Test filtering by name + response_by_name = get_python_environment(config, name="foo-package") + assert response_by_name.installed_packages == [ + PackageInfo( + name="foo-package", + version="http://example.com/foo.git @main", + location="", + is_dirty=False, + source=SourceInfo.SCRATCH, + ) + ] + + # Test filtering by source + response_by_source = get_python_environment(config, source=SourceInfo.SCRATCH) + assert response_by_source.installed_packages == [ + PackageInfo( + name="foo-package", + version="http://example.com/foo.git @main", + location="", + is_dirty=False, + source=SourceInfo.SCRATCH, + ) + ] + + +@pytest.fixture +def pyproject_file(tmp_path: Path) -> Generator[Path]: + pyproject_path = tmp_path / "pyproject.toml" + with pyproject_path.open("w") as f: + f.write( + """ + [project] + name = "example-project" + """ + ) + yield pyproject_path + os.remove(pyproject_path) + + +def test_get_project_name_from_pyproject_returns_name(pyproject_file: Path): + project_name = _get_project_name_from_pyproject(pyproject_file.parent) + assert project_name == "example-project" + + +def test_get_project_name_from_pyproject_returns_empty_if_no_pyproject( + tmp_path: Path, +): + project_name = _get_project_name_from_pyproject(tmp_path) + assert project_name == "" + + +def test_get_project_name_from_pyproject_returns_empty_if_no_name_key( + tmp_path: Path, +): + pyproject_path = tmp_path / "pyproject.toml" + with pyproject_path.open("w") as f: + f.write( + """ + [project] + version = "1.0.0" + """ + ) + project_name = _get_project_name_from_pyproject(tmp_path) + assert project_name == "" diff --git a/tests/unit_tests/service/test_interface.py b/tests/unit_tests/service/test_interface.py index 892c5ad2ea..0ddc69387f 100644 --- a/tests/unit_tests/service/test_interface.py +++ b/tests/unit_tests/service/test_interface.py @@ -37,8 +37,8 @@ TaskRequest, WorkerTask, ) +from blueapi.service.numtracker import NumtrackerClient from blueapi.utils.invalid_config_error import InvalidConfigError -from blueapi.utils.numtracker import NumtrackerClient from blueapi.worker.event import ( TaskResult, TaskStatus, @@ -499,7 +499,7 @@ def test_stomp_config_makes_no_client_when_disabled(mock_stomp_client: StompClie assert interface.stomp_client() is None -@patch("blueapi.cli.scratch._fetch_installed_packages_details") +@patch("blueapi.service.environment._fetch_installed_packages_details") def test_get_scratch_no_config(mock_fetch_installed_packages: Mock): interface.set_config(ApplicationConfig(scratch=None)) mock_fetch_installed_packages.return_value = [] @@ -545,7 +545,7 @@ def test_configure_numtracker(): assert nt._url.unicode_string() == "https://numtracker-example.com/graphql" -@patch("blueapi.utils.numtracker.httpx.AsyncClient.post") +@patch("blueapi.service.numtracker.httpx.AsyncClient.post") async def test_headers_are_cleared(mock_post): mock_response = Mock() mock_post.return_value = mock_response @@ -608,7 +608,7 @@ def test_numtracker_requires_instrument_metadata(): interface.set_config(ApplicationConfig()) -@patch("blueapi.utils.numtracker.NumtrackerClient.create_scan") +@patch("blueapi.service.numtracker.NumtrackerClient.create_scan") async def test_numtracker_create_scan_called_with_arguments_from_metadata( mock_create_scan, ): diff --git a/tests/unit_tests/utils/test_numtracker.py b/tests/unit_tests/service/test_numtracker.py similarity index 99% rename from tests/unit_tests/utils/test_numtracker.py rename to tests/unit_tests/service/test_numtracker.py index 1aa6df9e88..ee3546e242 100644 --- a/tests/unit_tests/utils/test_numtracker.py +++ b/tests/unit_tests/service/test_numtracker.py @@ -5,7 +5,7 @@ from pydantic import HttpUrl from pytest_httpx import HTTPXMock -from blueapi.utils.numtracker import ( +from blueapi.service.numtracker import ( DirectoryPath, NumtrackerClient, NumtrackerScanMutationResponse, diff --git a/tests/unit_tests/test_log.py b/tests/unit_tests/test_log.py index ba3325b151..f09eede82e 100644 --- a/tests/unit_tests/test_log.py +++ b/tests/unit_tests/test_log.py @@ -60,13 +60,13 @@ def mock_stream_handler_emit() -> Generator[Mock]: @pytest.fixture def mock_graylog_emit() -> Generator[Mock]: - with patch("blueapi.log.GELFTCPHandler.emit") as graylog_emit: + with patch("graypy.GELFTCPHandler.emit") as graylog_emit: yield graylog_emit MOCK_HANDLER_EMIT_STRINGS = [ "blueapi.log.logging.StreamHandler.emit", - "blueapi.log.GELFTCPHandler.emit", + "graypy.GELFTCPHandler.emit", ] diff --git a/uv.lock b/uv.lock index 8457c279a5..e3290a3135 100644 --- a/uv.lock +++ b/uv.lock @@ -449,29 +449,23 @@ wheels = [ name = "blueapi" source = { editable = "." } dependencies = [ - { name = "aioca" }, - { name = "aiohttp" }, { name = "bluesky", extra = ["plotting"] }, { name = "bluesky-stomp" }, { name = "click" }, { name = "event-model" }, - { name = "fastapi" }, { name = "gitpython" }, - { name = "graypy" }, - { name = "httpx" }, { name = "observability-utils" }, { name = "opentelemetry-distro" }, - { name = "opentelemetry-instrumentation-fastapi" }, { name = "ophyd-async" }, + { name = "packaging" }, { name = "pydantic" }, { name = "pydantic-settings" }, { name = "pyjwt", extra = ["crypto"] }, { name = "pyyaml" }, { name = "requests" }, { name = "stomp-py" }, - { name = "tiled", extra = ["client"] }, - { name = "tomlkit" }, - { name = "uvicorn" }, + { name = "tqdm" }, + { name = "websockets" }, ] [package.optional-dependencies] @@ -479,9 +473,21 @@ demo = [ { name = "dls-dodal" }, { name = "ophyd-async", extra = ["sim"] }, ] +server = [ + { name = "aioca" }, + { name = "aiohttp" }, + { name = "fastapi" }, + { name = "graypy" }, + { name = "httpx" }, + { name = "opentelemetry-instrumentation-fastapi" }, + { name = "tiled", extra = ["client"] }, + { name = "tomlkit" }, + { name = "uvicorn" }, +] [package.dev-dependencies] dev = [ + { name = "blueapi", extra = ["server"] }, { name = "copier" }, { name = "deepdiff" }, { name = "dls-dodal" }, @@ -518,36 +524,40 @@ dev = [ [package.metadata] requires-dist = [ - { name = "aioca" }, - { name = "aiohttp", specifier = ">=3.13.5" }, + { name = "aioca", marker = "extra == 'server'" }, + { name = "aiohttp", marker = "extra == 'server'", specifier = ">=3.13.5" }, { name = "bluesky", extras = ["plotting"], specifier = ">=1.14.0" }, { name = "bluesky-stomp", specifier = ">=0.1.6" }, { name = "click", specifier = ">=8.2.0" }, { name = "dls-dodal", marker = "extra == 'demo'", specifier = ">=1.69.0" }, { name = "event-model", specifier = "==1.24.0" }, - { name = "fastapi", specifier = ">=0.112.0" }, + { name = "fastapi", marker = "extra == 'server'", specifier = ">=0.112.0" }, { name = "gitpython", specifier = ">=3.1.58" }, - { name = "graypy", specifier = ">=2.1.0" }, - { name = "httpx", specifier = ">=0.28.1" }, + { name = "graypy", marker = "extra == 'server'", specifier = ">=2.1.0" }, + { name = "httpx", marker = "extra == 'server'", specifier = ">=0.28.1" }, { name = "observability-utils", specifier = ">=0.1.4" }, { name = "opentelemetry-distro", specifier = ">=0.48b0" }, - { name = "opentelemetry-instrumentation-fastapi", specifier = ">=0.48b0" }, + { name = "opentelemetry-instrumentation-fastapi", marker = "extra == 'server'", specifier = ">=0.48b0" }, { name = "ophyd-async", specifier = ">=0.13.5" }, { name = "ophyd-async", extras = ["sim"], marker = "extra == 'demo'" }, + { name = "packaging" }, { name = "pydantic", specifier = ">=2.0" }, { name = "pydantic-settings" }, { name = "pyjwt", extras = ["crypto"] }, { name = "pyyaml", specifier = ">=6.0.2" }, { name = "requests" }, { name = "stomp-py" }, - { name = "tiled", extras = ["client"], specifier = ">=0.2.4" }, - { name = "tomlkit" }, - { name = "uvicorn", specifier = ">=0.52.1" }, + { name = "tiled", extras = ["client"], marker = "extra == 'server'", specifier = ">=0.2.4" }, + { name = "tomlkit", marker = "extra == 'server'" }, + { name = "tqdm" }, + { name = "uvicorn", marker = "extra == 'server'", specifier = ">=0.52.1" }, + { name = "websockets" }, ] -provides-extras = ["demo"] +provides-extras = ["demo", "server"] [package.metadata.requires-dev] dev = [ + { name = "blueapi", extras = ["server"] }, { name = "copier" }, { name = "deepdiff" }, { name = "dls-dodal", specifier = ">=1.69.0" },