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:
@@ -303,6 +304,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, is_reusable=True)
     httpx_mock.add_callback(
@@ -331,9 +333,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(
@@ -350,10 +354,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(
@@ -130,16 +130,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:
@@ -272,6 +273,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, is_reusable=True)
     httpx_mock.add_callback(
@@ -294,9 +296,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(
@@ -307,10 +311,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(
@@ -82,6 +83,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"
@@ -112,9 +114,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(
@@ -125,10 +129,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")
@@ -213,6 +213,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,
@@ -307,6 +309,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()
@@ -323,9 +326,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,
@@ -337,9 +342,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(
@@ -352,10 +359,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, is_reusable=True)
-    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"