test(NoTicket): enable V2 auth refresh and invalidation coverage

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()

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()

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,

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,