unit tests

Author: ptiurin

tests/unit/V1/async_db/test_connection.py

@@ -12,14 +12,15 @@
 from firebolt.async_db.connection import Connection, connect
 from firebolt.client.auth import Auth, Token, UsernamePassword
 from firebolt.common._types import ColType
+from firebolt.utils.cache import _firebolt_cache
 from firebolt.utils.exception import (
     AccountNotFoundError,
     ConfigurationError,
     ConnectionClosedError,
     FireboltEngineError,
 )
-from firebolt.utils.token_storage import TokenSecureStorage
 from firebolt.utils.urls import ACCOUNT_ENGINE_ID_BY_NAME_URL
+from tests.unit.test_cache_helpers import get_cached_token
 
 
 async def test_closed_connection(connection: Connection) -> None:
@@ -275,6 +276,7 @@ async def test_connection_token_caching(
     access_token: str,
     account_id_callback: Callable,
     account_id_url: str,
+    enable_cache: Callable,
 ) -> None:
     httpx_mock.add_callback(check_credentials_callback, url=auth_url)
     httpx_mock.add_callback(query_callback, url=query_url)
@@ -295,9 +297,11 @@ async def test_connection_token_caching(
             assert await connection.cursor().execute("select*") == len(
                 python_query_data
             )
-        ts = TokenSecureStorage(username=user, password=password)
-        assert ts.get_cached_token() == access_token, "Invalid token value cached"
+        # Verify token was cached using the new cache system
+        cached_token = get_cached_token(user, password, account_name)
+        assert cached_token == access_token, "Invalid token value cached"
 
+    _firebolt_cache.clear()
     # Do the same, but with use_token_cache=False
     with Patcher():
         async with await connect(
@@ -314,10 +318,9 @@ async def test_connection_token_caching(
             assert await connection.cursor().execute("select*") == len(
                 python_query_data
             )
-        ts = TokenSecureStorage(username=user, password=password)
-        assert (
-            ts.get_cached_token() is None
-        ), "Token is cached even though caching is disabled"
+        # Verify token was not cached when caching is disabled
+        cached_token = get_cached_token(user, password, account_name)
+        assert cached_token is None, "Token is cached even though caching is disabled"
 
 
 async def test_connect_with_auth(

tests/unit/V1/client/test_client.py

@@ -10,9 +10,9 @@
 from firebolt.client import ClientV1 as Client
 from firebolt.client.auth import Token, UsernamePassword
 from firebolt.client.resource_manager_hooks import raise_on_4xx_5xx
-from firebolt.utils.token_storage import TokenSecureStorage
 from firebolt.utils.urls import AUTH_URL
 from firebolt.utils.util import fix_url_schema
+from tests.unit.test_cache_helpers import cache_token
 
 
 def test_client_retry(
@@ -121,16 +121,19 @@ def test_refresh_with_hooks(
     test_username: str,
     test_password: str,
     test_token: str,
+    enable_cache: Callable,
 ) -> 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)
+    cache_token(test_username, test_password, test_token, 2**32)
+    auth = UsernamePassword(test_username, test_password)
+    # Simulate what connect() would do
+    auth.account = None
 
     client = Client(
-        auth=UsernamePassword(test_username, test_password),
+        auth=auth,
         event_hooks={
             "response": [raise_on_4xx_5xx],
         },

tests/unit/V1/db/test_connection.py

@@ -13,14 +13,15 @@
 from firebolt.common._types import ColType
 from firebolt.db import Connection, connect
 from firebolt.db.cursor import CursorV1 as Cursor
+from firebolt.utils.cache import _firebolt_cache
 from firebolt.utils.exception import (
     AccountNotFoundError,
     ConfigurationError,
     ConnectionClosedError,
     FireboltEngineError,
 )
-from firebolt.utils.token_storage import TokenSecureStorage
 from firebolt.utils.urls import ACCOUNT_ENGINE_ID_BY_NAME_URL
+from tests.unit.test_cache_helpers import get_cached_token
 
 
 def test_closed_connection(connection: Connection) -> None:
@@ -248,6 +249,7 @@ def test_connection_token_caching(
     access_token: str,
     account_id_callback: Callable,
     account_id_url: str,
+    enable_cache: Callable,
 ) -> None:
     httpx_mock.add_callback(check_credentials_callback, url=auth_url)
     httpx_mock.add_callback(query_callback, url=query_url)
@@ -262,9 +264,11 @@ def test_connection_token_caching(
             api_endpoint=api_endpoint,
         ) as connection:
             assert connection.cursor().execute("select*") == len(python_query_data)
-        ts = TokenSecureStorage(username=user, password=password)
-        assert ts.get_cached_token() == access_token, "Invalid token value cached"
+        # Verify token was cached using the new cache system
+        cached_token = get_cached_token(user, password, account_name)
+        assert cached_token == access_token, "Invalid token value cached"
 
+    _firebolt_cache.clear()
     # Do the same, but with use_token_cache=False
     with Patcher():
         with connect(
@@ -275,10 +279,9 @@ def test_connection_token_caching(
             api_endpoint=api_endpoint,
         ) as connection:
             assert connection.cursor().execute("select*") == len(python_query_data)
-        ts = TokenSecureStorage(username=user, password=password)
-        assert (
-            ts.get_cached_token() is None
-        ), "Token is cached even though caching is disabled"
+        # Verify token was not cached when caching is disabled
+        cached_token = get_cached_token(user, password, account_name)
+        assert cached_token is None, "Token is cached even though caching is disabled"
 
 
 def test_connect_with_auth(

tests/unit/V1/service/test_resource_manager.py

@@ -8,8 +8,9 @@
 from firebolt.client.auth import Auth, Token, UsernamePassword
 from firebolt.common.settings import Settings
 from firebolt.service.manager import ResourceManager
+from firebolt.utils.cache import _firebolt_cache
 from firebolt.utils.exception import AccountNotFoundError
-from firebolt.utils.token_storage import TokenSecureStorage
+from tests.unit.test_cache_helpers import get_cached_token
 
 
 def test_rm_credentials(
@@ -70,6 +71,7 @@ def test_rm_token_cache(
     account_id_url: Pattern,
     account_id_callback: Callable,
     access_token: str,
+    enable_cache: Callable,
 ) -> None:
     """Credentials, that are passed to rm are processed properly."""
     url = "https://url"
@@ -88,9 +90,11 @@ def test_rm_token_cache(
         rm = ResourceManager(local_settings)
         rm._client.get(url)
 
-        ts = TokenSecureStorage(user, password)
-        assert ts.get_cached_token() == access_token, "Invalid token value cached"
+        # Verify token was cached using the new cache system
+        cached_token = get_cached_token(user, password)
+        assert cached_token == access_token, "Invalid token value cached"
 
+    _firebolt_cache.clear()
     # Do the same, but with use_token_cache=False
     with Patcher():
         local_settings = Settings(
@@ -101,10 +105,9 @@ def test_rm_token_cache(
         rm = ResourceManager(local_settings)
         rm._client.get(url)
 
-        ts = TokenSecureStorage(user, password)
-        assert (
-            ts.get_cached_token() is None
-        ), "Token is cached even though caching is disabled"
+        # Verify token was not cached when caching is disabled
+        cached_token = get_cached_token(user, password)
+        assert cached_token is None, "Token is cached even though caching is disabled"
 
 
 def test_rm_invalid_account_name(

tests/unit/async_db/test_connection.py

@@ -16,7 +16,7 @@
     ConnectionClosedError,
     FireboltError,
 )
-from firebolt.utils.token_storage import TokenSecureStorage
+from tests.unit.test_cache_helpers import get_cached_token
 
 
 @mark.skip("__slots__ is broken on Connection class")
@@ -202,6 +202,8 @@ def system_engine_callback_counter(request, **kwargs):
     else:
         assert system_engine_call_counter != 1, "System engine URL was cached"
 
+    _firebolt_cache.clear()
+
 
 async def test_connect_engine_failed(
     db_name: str,
@@ -290,6 +292,7 @@ async def test_connection_token_caching(
     python_query_data: List[List[ColType]],
     mock_connection_flow: Callable,
     mock_query: Callable,
+    enable_cache: Callable,
 ) -> None:
     mock_connection_flow()
     mock_query()
@@ -306,9 +309,11 @@ async def test_connection_token_caching(
             assert await connection.cursor().execute("select*") == len(
                 python_query_data
             )
-        ts = TokenSecureStorage(username=client_id, password=client_secret)
-        assert ts.get_cached_token() == access_token, "Invalid token value cached"
+        # Verify token was cached using the new cache system
+        cached_token = get_cached_token(client_id, client_secret, account_name)
+        assert cached_token == access_token, "Invalid token value cached"
 
+    _firebolt_cache.clear()
     with Patcher():
         async with await connect(
             database=db_name,
@@ -320,9 +325,11 @@ async def test_connection_token_caching(
             assert await connection.cursor().execute("select*") == len(
                 python_query_data
             )
-        ts = TokenSecureStorage(username=client_id, password=client_secret)
-        assert ts.get_cached_token() == access_token, "Invalid token value cached"
+        # Verify token was cached using the new cache system (second check)
+        cached_token = get_cached_token(client_id, client_secret, account_name)
+        assert cached_token == access_token, "Invalid token value cached"
 
+    _firebolt_cache.clear()
     # Do the same, but with use_token_cache=False
     with Patcher():
         async with await connect(
@@ -335,10 +342,9 @@ async def test_connection_token_caching(
             assert await connection.cursor().execute("select*") == len(
                 python_query_data
             )
-        ts = TokenSecureStorage(username=client_id, password=client_secret)
-        assert (
-            ts.get_cached_token() is None
-        ), "Token is cached even though caching is disabled"
+        # Verify token was not cached when caching is disabled
+        cached_token = get_cached_token(client_id, client_secret, account_name)
+        assert cached_token is None, "Token is cached even though caching is disabled"
 
 
 async def test_connect_with_user_agent(

tests/unit/client/auth/test_auth.py

@@ -1,13 +1,14 @@
 from types import MethodType
-from unittest.mock import PropertyMock, patch
+from typing import Generator
 
 from httpx import Request, codes
 from pyfakefs.fake_filesystem_unittest import Patcher
 from pytest import mark
 from pytest_httpx import HTTPXMock
 
-from firebolt.client.auth import Auth
-from firebolt.utils.token_storage import TokenSecureStorage
+from firebolt.client.auth import Auth, ClientCredentials
+from firebolt.utils.cache import _firebolt_cache
+from tests.unit.test_cache_helpers import get_cached_token
 from tests.unit.util import execute_generator_requests
 
 
@@ -89,6 +90,7 @@ def test_auth_token_storage(
     client_id: str,
     client_secret: str,
     access_token: str,
+    enable_cache: Generator,
 ) -> None:
     # Mock auth flow
     def set_token(token: str) -> callable:
@@ -101,29 +103,28 @@ def inner(self):
 
     url = "https://host"
     httpx_mock.add_response(status_code=codes.OK, url=url)
-    with Patcher(), patch(
-        "firebolt.client.auth.base.Auth._token_storage",
-        new_callable=PropertyMock,
-        return_value=TokenSecureStorage(client_id, client_secret),
-    ):
-        auth = Auth(use_token_cache=True)
+
+    # Test with caching enabled
+    with Patcher():
+        auth = ClientCredentials(client_id, client_secret, use_token_cache=True)
         # Get token
         auth.get_new_token_generator = MethodType(set_token(access_token), auth)
         execute_generator_requests(auth.auth_flow(Request("GET", url)))
 
-        st = TokenSecureStorage(client_id, client_secret)
-        assert st.get_cached_token() == access_token, "Invalid token value cached"
+        # Verify token was cached using the new cache system
+        cached_token = get_cached_token(client_id, client_secret, None)
+        assert cached_token == access_token, "Invalid token value cached"
+
+    # Clear cache before second test
+    _firebolt_cache.clear()
 
-    with Patcher(), patch(
-        "firebolt.client.auth.base.Auth._token_storage",
-        new_callable=PropertyMock,
-        return_value=TokenSecureStorage(client_id, client_secret),
-    ):
-        auth = Auth(use_token_cache=False)
+    # Test with caching disabled
+    with Patcher():
+        auth = ClientCredentials(client_id, client_secret, use_token_cache=False)
         # Get token
         auth.get_new_token_generator = MethodType(set_token(access_token), auth)
         execute_generator_requests(auth.auth_flow(Request("GET", url)))
-        st = TokenSecureStorage(client_id, client_secret)
-        assert (
-            st.get_cached_token() is None
-        ), "Token cached even though caching is disabled"
+
+        # Verify token was not cached
+        cached_token = get_cached_token(client_id, client_secret, None)
+        assert cached_token is None, "Token cached even though caching is disabled"