diff --git a/src/groundlight/client.py b/src/groundlight/client.py index 72cb74de..c988c8f1 100644 --- a/src/groundlight/client.py +++ b/src/groundlight/client.py @@ -125,6 +125,8 @@ class Groundlight: # pylint: disable=too-many-instance-attributes,too-many-publ Warning: Only disable verification when connecting to a Groundlight Edge Endpoint using self-signed certificates. For security, always keep verification enabled when using the Groundlight cloud service. + :param enable_token_rotation: If True (default), automatically rotate tokens whose identity has a + non-null Token TTL. :return: Groundlight client instance """ @@ -135,12 +137,13 @@ class Groundlight: # pylint: disable=too-many-instance-attributes,too-many-publ POLLING_EXPONENTIAL_BACKOFF = 1.3 # This still has the nice backoff property that the max number of requests # is O(log(time)), but with 1.3 the guarantee is that the call will return no more than 30% late - def __init__( + def __init__( # noqa: PLR0913 # pylint: disable=too-many-arguments self, endpoint: Optional[str] = None, api_token: Optional[str] = None, disable_tls_verification: Optional[bool] = None, http_transport_retries: Optional[Union[int, Retry]] = None, + enable_token_rotation: bool = True, ): """ Initialize a new Groundlight client instance. @@ -156,6 +159,8 @@ def __init__( certificates. For security, always keep verification enabled when using the Groundlight cloud service. :param http_transport_retries: Overrides urllib3 `PoolManager` retry policy for HTTP/HTTPS (forwarded to `Configuration.retries`). Not the same as SDK 5xx retries handled by `RequestsRetryDecorator`. + :param enable_token_rotation: If True (default), automatically rotate tokens whose identity has a + non-null Token TTL. :return: Groundlight client """ @@ -194,14 +199,16 @@ def __init__( self.configuration.api_key["ApiToken"] = api_token self.api_client = GroundlightApiClient(self.configuration) - try: - self._token_manager = TokenManager( - configured_token=api_token, - configuration=self.configuration, - request_timeout=DEFAULT_REQUEST_TIMEOUT, - ) - except TokenManagerError as exc: - raise ApiTokenError(str(exc)) from exc + self._token_manager: Optional[TokenManager] = None + if enable_token_rotation: + try: + self._token_manager = TokenManager( + configured_token=api_token, + configuration=self.configuration, + request_timeout=DEFAULT_REQUEST_TIMEOUT, + ) + except TokenManagerError as exc: + raise ApiTokenError(str(exc)) from exc self.detectors_api = DetectorsApi(self.api_client) self.detector_group_api = DetectorGroupsApi(self.api_client) self.images_api = ImageQueriesApi(self.api_client) @@ -212,7 +219,8 @@ def __init__( self.logged_in_user = "(not-logged-in)" self._verify_connectivity() # No-op when the working token has no identity Token TTL. - self._token_manager.start() + if self._token_manager is not None: + self._token_manager.start() def __repr__(self) -> str: # Don't call the API here because that can get us stuck in a loop rendering exception strings @@ -228,7 +236,8 @@ def __exit__(self, exc_type, exc_value, traceback) -> None: def close(self) -> None: """Stop the token refresh thread and close the HTTP client.""" - self._token_manager.close() + if self._token_manager is not None: + self._token_manager.close() self.api_client.close() def _verify_connectivity(self) -> None: diff --git a/src/groundlight/experimental_api.py b/src/groundlight/experimental_api.py index 8deb93e0..09258c6f 100644 --- a/src/groundlight/experimental_api.py +++ b/src/groundlight/experimental_api.py @@ -52,6 +52,7 @@ def __init__( endpoint: Union[str, None] = None, api_token: Union[str, None] = None, disable_tls_verification: Optional[bool] = None, + enable_token_rotation: bool = True, ): """ Constructs an experimental Groundlight client. @@ -83,8 +84,15 @@ def __init__( Warning: Only disable verification when connecting to a Groundlight Edge Endpoint using self-signed certificates. For security, always keep verification enabled when using the Groundlight cloud service. - """ - super().__init__(endpoint=endpoint, api_token=api_token, disable_tls_verification=disable_tls_verification) + :param enable_token_rotation: If True (default), automatically rotate tokens whose identity has a + non-null Token TTL. + """ + super().__init__( + endpoint=endpoint, + api_token=api_token, + disable_tls_verification=disable_tls_verification, + enable_token_rotation=enable_token_rotation, + ) self.notes_api = NotesApi(self.api_client) self.detector_group_api = DetectorGroupsApi(self.api_client) self.detector_reset_api = DetectorResetApi(self.api_client) diff --git a/test/unit/test_token_refresh_client.py b/test/unit/test_token_refresh_client.py index 7c58c0e2..07068415 100644 --- a/test/unit/test_token_refresh_client.py +++ b/test/unit/test_token_refresh_client.py @@ -18,3 +18,18 @@ def test_groundlight_starts_and_closes_token_manager(mocker): manager.close.assert_called_once() api_client_close.assert_called_once() + + +def test_groundlight_skips_token_manager_when_rotation_disabled(mocker): + """Disabling rotation leaves the configured token in place with no TokenManager.""" + token_manager_class = mocker.patch("groundlight.client.TokenManager") + mocker.patch.object(Groundlight, "_verify_connectivity") + + client = Groundlight(api_token="api_bootstrap_token_value_long_enough", enable_token_rotation=False) + api_client_close = mocker.patch.object(client.api_client, "close") + client.close() + + token_manager_class.assert_not_called() + assert client._token_manager is None + assert client.configuration.api_key["ApiToken"] == "api_bootstrap_token_value_long_enough" + api_client_close.assert_called_once()