Skip to content

Commit 589bdfe

Browse files
ptiurinptiurin
committed
additional tests
1 parent 5ee56bd commit 589bdfe

File tree

2 files changed

+676
-1
lines changed

2 files changed

+676
-1
lines changed

‎tests/unit/db/test_caching.py‎

Lines changed: 353 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import os
12
import time
23
from typing import Callable, Dict, Generator
34
from unittest.mock import patch
@@ -7,6 +8,7 @@
78
from pytest_httpx import HTTPXMock
89

910
from firebolt.client.auth import Auth
11+
from firebolt.client.auth.client_credentials import ClientCredentials
1012
from firebolt.db import connect
1113
from firebolt.utils.cache import CACHE_EXPIRY_SECONDS, _firebolt_cache
1214

@@ -512,7 +514,7 @@ def use_engine_callback_counter(request, **kwargs):
512514
)
513515

514516
# First connection - should populate cache
515-
connection_test(db_name, engine_name, True) # cache_enabled=True
517+
connection_test(db_name, engine_name, True)
516518

517519
# Verify initial calls were made
518520
assert system_engine_call_counter == 1, "System engine URL was not called initially"
@@ -629,6 +631,275 @@ def use_engine_callback_counter(request, **kwargs):
629631
assert use_engine_call_counter == 1, "USE ENGINE was called when cache should hit"
630632

631633

634+
def test_token_cache_expiry_forces_reauthentication(
635+
api_endpoint: str,
636+
auth: Auth,
637+
account_name: str,
638+
mock_connection_flow: Callable,
639+
):
640+
"""Test that expired authentication tokens force re-authentication."""
641+
mock_connection_flow()
642+
643+
fixed_time = 1000000
644+
645+
with patch("time.time", return_value=fixed_time):
646+
# First connection should cache token
647+
with connect(
648+
database="test_db",
649+
engine_name="test_engine",
650+
auth=auth,
651+
account_name=account_name,
652+
api_endpoint=api_endpoint,
653+
) as _:
654+
# Connection established and token cached
655+
assert auth.token is not None
656+
cached_token_1 = auth.token
657+
658+
# Simulate time passing within token validity
659+
with patch("time.time", return_value=fixed_time + 1800): # 30 minutes later
660+
with connect(
661+
database="test_db",
662+
engine_name="test_engine",
663+
auth=auth,
664+
account_name=account_name,
665+
api_endpoint=api_endpoint,
666+
) as _:
667+
# Should use cached token
668+
assert auth.token == cached_token_1
669+
670+
# Simulate time passing beyond token expiry
671+
expired_time = fixed_time + CACHE_EXPIRY_SECONDS + 100
672+
with patch("time.time", return_value=expired_time):
673+
# Clear the current token to simulate expiry
674+
auth._token = None
675+
auth._expires = fixed_time # Set to expired time
676+
677+
with connect(
678+
database="test_db",
679+
engine_name="test_engine",
680+
auth=auth,
681+
account_name=account_name,
682+
api_endpoint=api_endpoint,
683+
) as _:
684+
# Token expired, new authentication should occur
685+
assert auth.token is not None
686+
# Should get a new token (implementation detail may vary)
687+
assert auth.token is not None
688+
689+
690+
@mark.parametrize("use_cache", [True, False])
691+
def test_connection_cache_isolation_by_credentials(
692+
db_name: str,
693+
engine_name: str,
694+
auth_url: str,
695+
httpx_mock: HTTPXMock,
696+
check_credentials_callback: Callable,
697+
get_system_engine_url: str,
698+
get_system_engine_callback: Callable,
699+
system_engine_query_url: str,
700+
system_engine_no_db_query_url: str,
701+
query_url: str,
702+
use_database_callback: Callable,
703+
use_engine_callback: Callable,
704+
query_callback: Callable,
705+
account_name: str,
706+
api_endpoint: str,
707+
use_cache: bool,
708+
):
709+
"""Test that connections with different credentials are cached separately."""
710+
711+
# Create two different auth objects
712+
auth1 = ClientCredentials(
713+
client_id="client_1", client_secret="secret_1", use_token_cache=use_cache
714+
)
715+
auth2 = ClientCredentials(
716+
client_id="client_2", client_secret="secret_2", use_token_cache=use_cache
717+
)
718+
719+
system_engine_call_counter = 0
720+
721+
def system_engine_callback_counter(request, **kwargs):
722+
nonlocal system_engine_call_counter
723+
system_engine_call_counter += 1
724+
return get_system_engine_callback(request, **kwargs)
725+
726+
# Set up mocks for both auth credentials
727+
httpx_mock.add_callback(check_credentials_callback, url=auth_url, is_reusable=True)
728+
httpx_mock.add_callback(
729+
system_engine_callback_counter,
730+
url=get_system_engine_url,
731+
is_reusable=True,
732+
)
733+
httpx_mock.add_callback(
734+
use_database_callback,
735+
url=system_engine_no_db_query_url,
736+
match_content=f'USE DATABASE "{db_name}"'.encode("utf-8"),
737+
is_reusable=True,
738+
)
739+
httpx_mock.add_callback(
740+
use_engine_callback,
741+
url=system_engine_query_url,
742+
match_content=f'USE ENGINE "{engine_name}"'.encode("utf-8"),
743+
is_reusable=True,
744+
)
745+
httpx_mock.add_callback(
746+
query_callback,
747+
url=query_url,
748+
is_reusable=True,
749+
)
750+
751+
# Connect with first credentials
752+
with connect(
753+
database=db_name,
754+
engine_name=engine_name,
755+
auth=auth1,
756+
account_name=account_name,
757+
api_endpoint=api_endpoint,
758+
disable_cache=not use_cache,
759+
) as connection:
760+
connection.cursor().execute("SELECT 1")
761+
762+
first_call_count = system_engine_call_counter
763+
764+
# Connect with second credentials
765+
with connect(
766+
database=db_name,
767+
engine_name=engine_name,
768+
auth=auth2,
769+
account_name=account_name,
770+
api_endpoint=api_endpoint,
771+
disable_cache=not use_cache,
772+
) as connection:
773+
connection.cursor().execute("SELECT 1")
774+
775+
second_call_count = system_engine_call_counter
776+
777+
# Connect again with first credentials
778+
with connect(
779+
database=db_name,
780+
engine_name=engine_name,
781+
auth=auth1,
782+
account_name=account_name,
783+
api_endpoint=api_endpoint,
784+
disable_cache=not use_cache,
785+
) as connection:
786+
connection.cursor().execute("SELECT 1")
787+
788+
third_call_count = system_engine_call_counter
789+
790+
if use_cache:
791+
# With caching: each credential should only trigger system engine call once
792+
assert first_call_count == 1, "First auth should trigger system engine call"
793+
assert (
794+
second_call_count == 2
795+
), "Second auth should trigger another system engine call"
796+
assert third_call_count == 2, "First auth second use should use cache"
797+
else:
798+
# Without caching: every connection should trigger system engine call
799+
assert first_call_count >= 1, "System engine should be called"
800+
assert (
801+
second_call_count > first_call_count
802+
), "Each connection should call system engine"
803+
assert third_call_count > second_call_count, "No caching means more calls"
804+
805+
806+
def test_connection_cache_isolation_by_accounts(
807+
db_name: str,
808+
engine_name: str,
809+
auth_url: str,
810+
httpx_mock: HTTPXMock,
811+
check_credentials_callback: Callable,
812+
get_system_engine_url: str,
813+
get_system_engine_callback: Callable,
814+
system_engine_query_url: str,
815+
system_engine_no_db_query_url: str,
816+
query_url: str,
817+
use_database_callback: Callable,
818+
use_engine_callback: Callable,
819+
query_callback: Callable,
820+
auth: Auth,
821+
api_endpoint: str,
822+
):
823+
"""Test that connections to different accounts are cached separately."""
824+
system_engine_call_counter = 0
825+
826+
def system_engine_callback_counter(request, **kwargs):
827+
nonlocal system_engine_call_counter
828+
system_engine_call_counter += 1
829+
return get_system_engine_callback(request, **kwargs)
830+
831+
httpx_mock.add_callback(check_credentials_callback, url=auth_url, is_reusable=True)
832+
httpx_mock.add_callback(
833+
system_engine_callback_counter,
834+
url=get_system_engine_url,
835+
is_reusable=True,
836+
)
837+
httpx_mock.add_callback(
838+
use_database_callback,
839+
url=system_engine_no_db_query_url,
840+
match_content=f'USE DATABASE "{db_name}"'.encode("utf-8"),
841+
is_reusable=True,
842+
)
843+
httpx_mock.add_callback(
844+
use_engine_callback,
845+
url=system_engine_query_url,
846+
match_content=f'USE ENGINE "{engine_name}"'.encode("utf-8"),
847+
is_reusable=True,
848+
)
849+
httpx_mock.add_callback(
850+
query_callback,
851+
url=query_url,
852+
is_reusable=True,
853+
)
854+
855+
# Connect to first account
856+
with connect(
857+
database=db_name,
858+
engine_name=engine_name,
859+
auth=auth,
860+
account_name="account_1",
861+
api_endpoint=api_endpoint,
862+
) as connection:
863+
connection.cursor().execute("SELECT 1")
864+
865+
first_account_calls = system_engine_call_counter
866+
867+
# Connect to second account with same credentials
868+
with connect(
869+
database=db_name,
870+
engine_name=engine_name,
871+
auth=auth,
872+
account_name="account_2",
873+
api_endpoint=api_endpoint,
874+
) as connection:
875+
connection.cursor().execute("SELECT 1")
876+
877+
second_account_calls = system_engine_call_counter
878+
879+
# Connect again to first account
880+
with connect(
881+
database=db_name,
882+
engine_name=engine_name,
883+
auth=auth,
884+
account_name="account_1",
885+
api_endpoint=api_endpoint,
886+
) as connection:
887+
connection.cursor().execute("SELECT 1")
888+
889+
third_connection_calls = system_engine_call_counter
890+
891+
# Each account should require its own system engine call initially
892+
assert first_account_calls == 1, "First account should trigger system engine call"
893+
assert (
894+
second_account_calls == 2
895+
), "Second account should trigger another system engine call"
896+
897+
# Return to first account should use cached data
898+
assert (
899+
third_connection_calls == 2
900+
), "First account second connection should use cache"
901+
902+
632903
def test_database_switching_with_same_engine_preserves_database_context(
633904
db_name: str,
634905
engine_name: str,
@@ -874,3 +1145,84 @@ def flexible_query_callback(request, **kwargs):
8741145
f"Engine parameters should be separate from database context. "
8751146
f"Full engine params: {cached_engine.params}"
8761147
)
1148+
1149+
1150+
@mark.parametrize("cache_disabled", [True, False])
1151+
def test_cache_disable_via_environment_integration(
1152+
db_name: str,
1153+
engine_name: str,
1154+
auth_url: str,
1155+
httpx_mock: HTTPXMock,
1156+
check_credentials_callback: Callable,
1157+
get_system_engine_url: str,
1158+
get_system_engine_callback: Callable,
1159+
system_engine_query_url: str,
1160+
system_engine_no_db_query_url: str,
1161+
query_url: str,
1162+
use_database_callback: Callable,
1163+
use_engine_callback: Callable,
1164+
query_callback: Callable,
1165+
auth: Auth,
1166+
account_name: str,
1167+
api_endpoint: str,
1168+
cache_disabled: bool,
1169+
):
1170+
"""Test cache behavior when disabled via environment variable."""
1171+
system_engine_call_counter = 0
1172+
1173+
def system_engine_callback_counter(request, **kwargs):
1174+
nonlocal system_engine_call_counter
1175+
system_engine_call_counter += 1
1176+
return get_system_engine_callback(request, **kwargs)
1177+
1178+
httpx_mock.add_callback(check_credentials_callback, url=auth_url, is_reusable=True)
1179+
httpx_mock.add_callback(
1180+
system_engine_callback_counter,
1181+
url=get_system_engine_url,
1182+
is_reusable=True,
1183+
)
1184+
httpx_mock.add_callback(
1185+
use_database_callback,
1186+
url=system_engine_no_db_query_url,
1187+
match_content=f'USE DATABASE "{db_name}"'.encode("utf-8"),
1188+
is_reusable=True,
1189+
)
1190+
httpx_mock.add_callback(
1191+
use_engine_callback,
1192+
url=system_engine_query_url,
1193+
match_content=f'USE ENGINE "{engine_name}"'.encode("utf-8"),
1194+
is_reusable=True,
1195+
)
1196+
httpx_mock.add_callback(
1197+
query_callback,
1198+
url=query_url,
1199+
is_reusable=True,
1200+
)
1201+
1202+
# Set environment variable to disable cache
1203+
env_patch = {}
1204+
if cache_disabled:
1205+
env_patch["FIREBOLT_SDK_DISABLE_CACHE"] = "true"
1206+
1207+
with patch.dict(os.environ, env_patch):
1208+
# Create multiple connections to test caching behavior
1209+
for _ in range(3):
1210+
with connect(
1211+
database=db_name,
1212+
engine_name=engine_name,
1213+
auth=auth,
1214+
account_name=account_name,
1215+
api_endpoint=api_endpoint,
1216+
) as connection:
1217+
connection.cursor().execute("SELECT 1")
1218+
1219+
if cache_disabled:
1220+
# Each connection should call system engine when cache is disabled
1221+
assert (
1222+
system_engine_call_counter == 3
1223+
), "Cache disabled should not reuse system engine"
1224+
else:
1225+
# Only first connection should call system engine when cache is enabled
1226+
assert (
1227+
system_engine_call_counter == 1
1228+
), "Cache enabled should reuse system engine"

0 commit comments

Comments
 (0)