diff --git a/src/firebolt/client/resource_manager_hooks.py b/src/firebolt/client/resource_manager_hooks.py index 6cbbc99d9c6..aec44bbdfd3 100644 --- a/src/firebolt/client/resource_manager_hooks.py +++ b/src/firebolt/client/resource_manager_hooks.py @@ -1,7 +1,7 @@ from json import JSONDecodeError from logging import getLogger -from httpx import HTTPStatusError, Request, RequestError, Response +from httpx import HTTPStatusError, Request, RequestError, Response, codes logger = getLogger(__name__) @@ -51,6 +51,9 @@ def raise_on_4xx_5xx(response: Response) -> None: HTTPStatusError: HTTP error """ try: + if response.status_code == codes.UNAUTHORIZED: + # This should be handled by Auth + return response.raise_for_status() except RequestError as exc: logger.debug("An error occurred while requesting %s.", exc.request.url) diff --git a/tests/unit/client/test_client.py b/tests/unit/client/test_client.py index ea69ebbd8f4..aba06631178 100644 --- a/tests/unit/client/test_client.py +++ b/tests/unit/client/test_client.py @@ -2,12 +2,15 @@ from typing import Callable from httpx import codes +from pyfakefs.fake_filesystem import FakeFilesystem from pytest import raises from pytest_httpx import HTTPXMock from firebolt.client import DEFAULT_API_URL, Client from firebolt.client.auth import Token, UsernamePassword +from firebolt.client.resource_manager_hooks import raise_on_4xx_5xx from firebolt.common import Settings +from firebolt.utils.token_storage import TokenSecureStorage from firebolt.utils.urls import AUTH_URL from firebolt.utils.util import fix_url_schema @@ -109,3 +112,46 @@ def test_client_account_id( api_endpoint=settings.server, ) as c: assert c.account_id == account_id, "Invalid account id returned" + + +# FIR-14945 +def test_refresh_with_hooks( + fs: FakeFilesystem, + httpx_mock: HTTPXMock, + test_username: str, + test_password: str, + test_token: str, +) -> None: + """ + When hooks are used, the invalid token, fetched from cache, is refreshed + """ + + tss = TokenSecureStorage(test_username, test_password) + tss.cache_token(test_token, 2**32) + + client = Client( + auth=UsernamePassword(test_username, test_password), + event_hooks={ + "response": [raise_on_4xx_5xx], + }, + ) + + # client request failed authorization + httpx_mock.add_response( + status_code=codes.UNAUTHORIZED, + ) + + # auth get another token + httpx_mock.add_response( + status_code=codes.OK, + json={"expires_in": 2**30, "access_token": test_token}, + ) + + # client request success this time + httpx_mock.add_response( + status_code=codes.OK, + ) + + assert ( + client.get("https://url").status_code == codes.OK + ), "request failed with firebolt client"