diff --git a/tests/integration/dbapi/async/V2/test_auth_async.py b/tests/integration/dbapi/async/V2/test_auth_async.py index 98f71a2dcba..4989304a927 100644 --- a/tests/integration/dbapi/async/V2/test_auth_async.py +++ b/tests/integration/dbapi/async/V2/test_auth_async.py @@ -1,54 +1,54 @@ -from time import time - -from pytest import mark, raises +from pytest import fixture, mark, raises +from pytest_mock import MockerFixture from firebolt.async_db import Connection -from firebolt.utils.exception import AuthenticationError +from firebolt.client.auth import ClientCredentials +from firebolt.utils.exception import AuthenticationError, AuthorizationError +from tests.integration.conftest import Secret +pytestmark = mark.parametrize("connection_factory", ["remote"], indirect=True) -@mark.skip(reason="flaky, token not updated each time") -async def test_refresh_token(connection: Connection) -> None: - """Auth refreshes token on expiration/invalidation""" - async with connection.cursor() as c: - # Works fine - await c.execute("show tables") - # Invalidate the token - c._client.auth._token += "_" +@fixture +def auth(service_id: str, service_secret: Secret) -> ClientCredentials: + # These tests mutate credentials; never reuse the session auth or its cache. + return ClientCredentials(service_id, service_secret.value, use_token_cache=False) - # Still works fine - await c.execute("show tables") - old = c._client.auth.token - c._client.auth._expires = int(time()) - 1 +@mark.parametrize("invalidation", ["expired", "invalid"]) +async def test_refresh_token( + connection: Connection, mocker: MockerFixture, invalidation: str +) -> None: + async with connection.cursor() as c: + await c.execute("SELECT 1") + assert await c.fetchone() == [1] - # Still works fine. - await c.execute("show tables") + auth = c._client.auth + refresh = mocker.spy(auth, "get_new_token_generator") + if invalidation == "expired": + auth._expires = 0 + else: + auth._token = "invalid-token" - assert c._client.auth.token != old, "Auth didn't update token on expiration." + await c.execute("SELECT 1") + assert await c.fetchone() == [1] + refresh.assert_called_once() + assert auth.token + assert not auth.expired -@mark.skip("Avoiding excessive load with username/password") async def test_credentials_invalidation( - connection: Connection, username_password_connection: Connection + connection: Connection, mocker: MockerFixture ) -> None: - """Auth raises authentication error on credentials invalidation""" - # Can't pytest.parametrize it due to nested event loop error - for conn in [connection, username_password_connection]: - async with conn.cursor() as c: - # Works fine - await c.execute("show tables") - - # Invalidate the token - c._client.auth._token += "_" - # Invalidate credentials - for cred in ("username", "password", "client_id", "client_secret"): - if hasattr(c._client.auth, cred): - setattr(c._client.auth, cred, "_") - - with raises(AuthenticationError) as exc_info: - await c.execute("show tables") - - assert str(exc_info.value).startswith( - "Failed to authenticate" - ), "Invalid authentication error message" + async with connection.cursor() as c: + await c.execute("SELECT 1") + assert await c.fetchone() == [1] + + auth = c._client.auth + refresh = mocker.spy(auth, "get_new_token_generator") + auth._token = "invalid-token" + auth.client_secret = "_" + + with raises((AuthenticationError, AuthorizationError)): + await c.execute("SELECT 1") + refresh.assert_called_once() diff --git a/tests/integration/dbapi/sync/V2/test_auth.py b/tests/integration/dbapi/sync/V2/test_auth.py index 1e62a359a30..83480a6410f 100644 --- a/tests/integration/dbapi/sync/V2/test_auth.py +++ b/tests/integration/dbapi/sync/V2/test_auth.py @@ -1,51 +1,54 @@ -from time import time - -from pytest import mark, raises +from pytest import fixture, mark, raises +from pytest_mock import MockerFixture +from firebolt.client.auth import ClientCredentials from firebolt.db import Connection -from firebolt.utils.exception import AuthenticationError - - -@mark.skip(reason="flaky, token not updated each time") -def test_refresh_token(connection: Connection) -> None: - """Auth refreshes token on expiration/invalidation""" - with connection.cursor() as c: - # Works fine - c.execute("show tables") - - # Invalidate the token - c._client.auth._token += "_" +from firebolt.utils.exception import AuthenticationError, AuthorizationError +from tests.integration.conftest import Secret - # Still works fine - c.execute("show tables") +pytestmark = mark.parametrize("connection_factory", ["remote"], indirect=True) - old = c._client.auth.token - c._client.auth._expires = int(time()) - 1 - # Still works fine - c.execute("show tables") +@fixture +def auth(service_id: str, service_secret: Secret) -> ClientCredentials: + # These tests mutate credentials; never reuse the session auth or its cache. + return ClientCredentials(service_id, service_secret.value, use_token_cache=False) - assert c._client.auth.token != old, "Auth didn't update token on expiration" +@mark.parametrize("invalidation", ["expired", "invalid"]) +def test_refresh_token( + connection: Connection, mocker: MockerFixture, invalidation: str +) -> None: + with connection.cursor() as c: + c.execute("SELECT 1") + assert c.fetchone() == [1] + + auth = c._client.auth + refresh = mocker.spy(auth, "get_new_token_generator") + if invalidation == "expired": + auth._expires = 0 + else: + auth._token = "invalid-token" + + c.execute("SELECT 1") + assert c.fetchone() == [1] + refresh.assert_called_once() + assert auth.token + assert not auth.expired + + +def test_credentials_invalidation( + connection: Connection, mocker: MockerFixture +) -> None: + with connection.cursor() as c: + c.execute("SELECT 1") + assert c.fetchone() == [1] -@mark.skip("Avoiding excessive load with username/password") -@mark.parametrize("connection_fixture", ["connection", "username_password_connection"]) -def test_credentials_invalidation(connection_fixture: str, request) -> None: - """Auth raises Authentication Error on credentials invalidation""" - with request.getfixturevalue(connection_fixture).cursor() as c: - # Works fine - c.execute("show tables") - - # Invalidate the token - c._client.auth._token += "_" - # Invalidate credentials - for cred in ("username", "password", "client_id", "client_secret"): - if hasattr(c._client.auth, cred): - setattr(c._client.auth, cred, "_") - - with raises(AuthenticationError) as exc_info: - c.execute("show tables") + auth = c._client.auth + refresh = mocker.spy(auth, "get_new_token_generator") + auth._token = "invalid-token" + auth.client_secret = "_" - assert str(exc_info.value).startswith( - "Failed to authenticate" - ), "Invalid authentication error message" + with raises((AuthenticationError, AuthorizationError)): + c.execute("SELECT 1") + refresh.assert_called_once() diff --git a/tests/unit/client/test_client.py b/tests/unit/client/test_client.py index 1900151c67d..fd98e088574 100644 --- a/tests/unit/client/test_client.py +++ b/tests/unit/client/test_client.py @@ -2,12 +2,13 @@ from httpx import Request, Timeout, codes from pyfakefs.fake_filesystem import FakeFilesystem -from pytest import raises +from pytest import mark, raises from pytest_httpx import HTTPXMock from firebolt.client import ClientV2 as Client from firebolt.client.auth import Auth, ClientCredentials from firebolt.client.resource_manager_hooks import raise_on_4xx_5xx +from firebolt.utils.exception import AuthenticationError, AuthorizationError from firebolt.utils.token_storage import TokenSecureStorage from firebolt.utils.urls import AUTH_SERVICE_ACCOUNT_URL from tests.unit.conftest import Response @@ -51,6 +52,78 @@ def test_client_retry( ), "request failed with firebolt client" +def test_client_refresh_expired_token_with_same_value( + httpx_mock: HTTPXMock, + auth: Auth, + account_name: str, + api_endpoint: str, + auth_url: str, + access_token: str, +) -> None: + auth._token = access_token + auth._expires = 2**32 + httpx_mock.add_response( + url="https://url", + match_headers={"Authorization": f"Bearer {access_token}"}, + is_reusable=True, + ) + httpx_mock.add_response( + url=auth_url, json={"expires_in": 3600, "access_token": access_token} + ) + + with Client( + account_name=account_name, auth=auth, api_endpoint=api_endpoint + ) as client: + assert (client.get("https://url")).status_code == codes.OK + auth._expires = 0 + assert (client.get("https://url")).status_code == codes.OK + + assert auth.token == access_token + assert not auth.expired + assert len(httpx_mock.get_requests(url=auth_url)) == 1 + + +@mark.parametrize( + "status_code, error", + [ + (codes.BAD_REQUEST, AuthenticationError), + (codes.UNAUTHORIZED, AuthorizationError), + ], +) +def test_client_refresh_with_invalid_credentials( + httpx_mock: HTTPXMock, + auth: ClientCredentials, + account_name: str, + api_endpoint: str, + auth_url: str, + access_token: str, + status_code: int, + error: type[Exception], +) -> None: + auth._token = access_token + auth._expires = 2**32 + httpx_mock.add_response( + url="https://url", match_headers={"Authorization": f"Bearer {access_token}"} + ) + httpx_mock.add_response( + url="https://url", + status_code=codes.UNAUTHORIZED, + match_headers={"Authorization": "Bearer invalid-token"}, + ) + httpx_mock.add_response(url=auth_url, status_code=status_code) + + with Client( + account_name=account_name, auth=auth, api_endpoint=api_endpoint + ) as client: + assert (client.get("https://url")).status_code == codes.OK + auth._token = "invalid-token" + auth.client_secret = "_" + with raises(error): + client.get("https://url") + + assert len(httpx_mock.get_requests(url=auth_url)) == 1 + + def test_client_different_auths( httpx_mock: HTTPXMock, check_credentials_callback: Callable, diff --git a/tests/unit/client/test_client_async.py b/tests/unit/client/test_client_async.py index e81fb567920..92985a7aa78 100644 --- a/tests/unit/client/test_client_async.py +++ b/tests/unit/client/test_client_async.py @@ -5,12 +5,13 @@ from typing import Any, Callable from httpx import Request, Timeout, codes -from pytest import raises +from pytest import mark, raises from pytest_httpx import HTTPXMock from trio import open_nursery, sleep from firebolt.client import AsyncClientV2 as AsyncClient from firebolt.client.auth import Auth, ClientCredentials +from firebolt.utils.exception import AuthenticationError, AuthorizationError from firebolt.utils.urls import AUTH_SERVICE_ACCOUNT_URL from tests.unit.conftest import Response, retry_if_failed @@ -53,6 +54,78 @@ async def test_client_retry( ).status_code == codes.OK, "request failed with firebolt client" +async def test_client_refresh_expired_token_with_same_value( + httpx_mock: HTTPXMock, + auth: Auth, + account_name: str, + api_endpoint: str, + auth_url: str, + access_token: str, +) -> None: + auth._token = access_token + auth._expires = 2**32 + httpx_mock.add_response( + url="https://url", + match_headers={"Authorization": f"Bearer {access_token}"}, + is_reusable=True, + ) + httpx_mock.add_response( + url=auth_url, json={"expires_in": 3600, "access_token": access_token} + ) + + async with AsyncClient( + account_name=account_name, auth=auth, api_endpoint=api_endpoint + ) as client: + assert (await client.get("https://url")).status_code == codes.OK + auth._expires = 0 + assert (await client.get("https://url")).status_code == codes.OK + + assert auth.token == access_token + assert not auth.expired + assert len(httpx_mock.get_requests(url=auth_url)) == 1 + + +@mark.parametrize( + "status_code, error", + [ + (codes.BAD_REQUEST, AuthenticationError), + (codes.UNAUTHORIZED, AuthorizationError), + ], +) +async def test_client_refresh_with_invalid_credentials( + httpx_mock: HTTPXMock, + auth: ClientCredentials, + account_name: str, + api_endpoint: str, + auth_url: str, + access_token: str, + status_code: int, + error: type[Exception], +) -> None: + auth._token = access_token + auth._expires = 2**32 + httpx_mock.add_response( + url="https://url", match_headers={"Authorization": f"Bearer {access_token}"} + ) + httpx_mock.add_response( + url="https://url", + status_code=codes.UNAUTHORIZED, + match_headers={"Authorization": "Bearer invalid-token"}, + ) + httpx_mock.add_response(url=auth_url, status_code=status_code) + + async with AsyncClient( + account_name=account_name, auth=auth, api_endpoint=api_endpoint + ) as client: + assert (await client.get("https://url")).status_code == codes.OK + auth._token = "invalid-token" + auth.client_secret = "_" + with raises(error): + await client.get("https://url") + + assert len(httpx_mock.get_requests(url=auth_url)) == 1 + + async def test_client_different_auths( httpx_mock: HTTPXMock, check_credentials_callback: Callable,