More refactoring

Author: tisnik

tests/unit/authentication/test_k8s.py

@@ -5,6 +5,8 @@
 import os
 
 import pytest
+from pytest_mock import MockerFixture
+
 from fastapi import HTTPException, Request
 from kubernetes.client import AuthenticationV1Api, AuthorizationV1Api
 from kubernetes.client.rest import ApiException
@@ -70,7 +72,7 @@ def test_singleton_pattern():
     assert k1 is k2
 
 
-async def test_auth_dependency_valid_token(mocker):
+async def test_auth_dependency_valid_token(mocker: MockerFixture):
     """Tests the auth dependency with a mocked valid-token."""
     dependency = K8SAuthDependency()
 
@@ -103,7 +105,7 @@ async def test_auth_dependency_valid_token(mocker):
     assert token == "valid-token"
 
 
-async def test_auth_dependency_invalid_token(mocker):
+async def test_auth_dependency_invalid_token(mocker: MockerFixture):
     """Test the auth dependency with a mocked invalid-token."""
     dependency = K8SAuthDependency()
 
@@ -135,7 +137,7 @@ async def test_auth_dependency_invalid_token(mocker):
     assert exc_info.value.status_code == 403
 
 
-async def test_cluster_id_is_used_for_kube_admin(mocker):
+async def test_cluster_id_is_used_for_kube_admin(mocker: MockerFixture):
     """Test the cluster id is used as user_id when user is kube:admin."""
     dependency = K8SAuthDependency()
     mock_authz_api = mocker.patch("authentication.k8s.K8sClientSingleton.get_authz_api")
@@ -175,7 +177,7 @@ async def test_cluster_id_is_used_for_kube_admin(mocker):
     assert token == "valid-token"
 
 
-def test_auth_dependency_config(mocker):
+def test_auth_dependency_config(mocker: MockerFixture):
     """Test the auth dependency can load kubeconfig file."""
     mocker.patch.dict(os.environ, {"MY_ENV_VAR": "mocked"})
 
@@ -189,7 +191,7 @@ def test_auth_dependency_config(mocker):
     ), "authz_client is not an instance of AuthorizationV1Api"
 
 
-def test_get_cluster_id(mocker):
+def test_get_cluster_id(mocker: MockerFixture):
     """Test get_cluster_id function."""
     mock_get_custom_objects_api = mocker.patch(
         "authentication.k8s.K8sClientSingleton.get_custom_objects_api"
@@ -228,7 +230,7 @@ def test_get_cluster_id(mocker):
         K8sClientSingleton._get_cluster_id()
 
 
-def test_get_cluster_id_in_cluster(mocker):
+def test_get_cluster_id_in_cluster(mocker: MockerFixture):
     """Test get_cluster_id function when running inside of cluster."""
     mocker.patch("authentication.k8s.RUNNING_IN_CLUSTER", True)
     mocker.patch("authentication.k8s.K8sClientSingleton.__new__")
@@ -240,7 +242,7 @@ def test_get_cluster_id_in_cluster(mocker):
     assert K8sClientSingleton.get_cluster_id() == "some-cluster-id"
 
 
-def test_get_cluster_id_outside_of_cluster(mocker):
+def test_get_cluster_id_outside_of_cluster(mocker: MockerFixture):
     """Test get_cluster_id function when running outside of cluster."""
     mocker.patch("authentication.k8s.RUNNING_IN_CLUSTER", False)
     mocker.patch("authentication.k8s.K8sClientSingleton.__new__")

tests/unit/cache/test_cache_factory.py

@@ -1,6 +1,7 @@
 """Unit tests for CacheFactory class."""
 
 import pytest
+from pytest_mock import MockerFixture
 
 from constants import (
     CACHE_TYPE_NOOP,
@@ -112,7 +113,9 @@ def test_conversation_cache_sqlite_improper_config(tmpdir):
         _ = CacheFactory.conversation_cache(cc)
 
 
-def test_conversation_cache_postgres(postgres_cache_config_fixture, mocker):
+def test_conversation_cache_postgres(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Check if PostgreSQL is returned by factory with proper configuration."""
     mocker.patch("psycopg2.connect")
     cache = CacheFactory.conversation_cache(postgres_cache_config_fixture)

tests/unit/cache/test_postgres_cache.py

@@ -1,6 +1,7 @@
 """Unit tests for PostgreSQL cache implementation."""
 
 import pytest
+from pytest_mock import MockerFixture
 
 import psycopg2
 
@@ -69,7 +70,7 @@ def postgres_cache_config():
     )
 
 
-def test_cache_initialization(postgres_cache_config_fixture, mocker):
+def test_cache_initialization(postgres_cache_config_fixture, mocker: MockerFixture):
     """Test the get operation when DB is connected."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -80,7 +81,9 @@ def test_cache_initialization(postgres_cache_config_fixture, mocker):
     assert cache.connection is not None
 
 
-def test_cache_initialization_on_error(postgres_cache_config_fixture, mocker):
+def test_cache_initialization_on_error(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test the get operation when DB is not connected."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect", side_effect=Exception("foo"))
@@ -90,7 +93,9 @@ def test_cache_initialization_on_error(postgres_cache_config_fixture, mocker):
         _ = PostgresCache(postgres_cache_config_fixture)
 
 
-def test_cache_initialization_connect_finalizer(postgres_cache_config_fixture, mocker):
+def test_cache_initialization_connect_finalizer(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test the get operation when DB is not connected."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -106,7 +111,7 @@ def test_cache_initialization_connect_finalizer(postgres_cache_config_fixture, m
         _ = PostgresCache(postgres_cache_config_fixture)
 
 
-def test_connected_when_connected(postgres_cache_config_fixture, mocker):
+def test_connected_when_connected(postgres_cache_config_fixture, mocker: MockerFixture):
     """Test the connected() method."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -116,7 +121,9 @@ def test_connected_when_connected(postgres_cache_config_fixture, mocker):
     assert cache.connected() is True
 
 
-def test_connected_when_disconnected(postgres_cache_config_fixture, mocker):
+def test_connected_when_disconnected(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test the connected() method."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -128,7 +135,9 @@ def test_connected_when_disconnected(postgres_cache_config_fixture, mocker):
     assert cache.connected() is False
 
 
-def test_connected_when_connection_error(postgres_cache_config_fixture, mocker):
+def test_connected_when_connection_error(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test the connected() method."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -139,7 +148,9 @@ def test_connected_when_connection_error(postgres_cache_config_fixture, mocker):
     assert cache.connected() is False
 
 
-def test_initialize_cache_when_connected(postgres_cache_config_fixture, mocker):
+def test_initialize_cache_when_connected(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test the initialize_cache()."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -148,7 +159,9 @@ def test_initialize_cache_when_connected(postgres_cache_config_fixture, mocker):
     cache.initialize_cache()
 
 
-def test_initialize_cache_when_disconnected(postgres_cache_config_fixture, mocker):
+def test_initialize_cache_when_disconnected(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test the initialize_cache()."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -159,7 +172,7 @@ def test_initialize_cache_when_disconnected(postgres_cache_config_fixture, mocke
         cache.initialize_cache()
 
 
-def test_ready_method(postgres_cache_config_fixture, mocker):
+def test_ready_method(postgres_cache_config_fixture, mocker: MockerFixture):
     """Test the ready() method."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -170,7 +183,9 @@ def test_ready_method(postgres_cache_config_fixture, mocker):
     assert ready is True
 
 
-def test_get_operation_when_disconnected(postgres_cache_config_fixture, mocker):
+def test_get_operation_when_disconnected(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test the get() method."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -184,7 +199,9 @@ def test_get_operation_when_disconnected(postgres_cache_config_fixture, mocker):
         cache.get(USER_ID_1, CONVERSATION_ID_1, False)
 
 
-def test_get_operation_when_connected(postgres_cache_config_fixture, mocker):
+def test_get_operation_when_connected(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test the get() method."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -203,7 +220,9 @@ def test_get_operation_returned_values():
     # Need to mock the cursor.execute() method
 
 
-def test_insert_or_append_when_disconnected(postgres_cache_config_fixture, mocker):
+def test_insert_or_append_when_disconnected(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test the insert_or_append() method."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -217,7 +236,7 @@ def test_insert_or_append_when_disconnected(postgres_cache_config_fixture, mocke
 
 
 def test_insert_or_append_operation_when_connected(
-    postgres_cache_config_fixture, mocker
+    postgres_cache_config_fixture, mocker: MockerFixture
 ):
     """Test the insert_or_append() method."""
     # prevent real connection to PG instance
@@ -229,7 +248,7 @@ def test_insert_or_append_operation_when_connected(
 
 
 def test_insert_or_append_operation_operation_error(
-    postgres_cache_config_fixture, mocker
+    postgres_cache_config_fixture, mocker: MockerFixture
 ):
     """Test the insert_or_append() method."""
     # prevent real connection to PG instance
@@ -244,7 +263,7 @@ def test_insert_or_append_operation_operation_error(
         cache.insert_or_append(USER_ID_1, CONVERSATION_ID_1, cache_entry_1, False)
 
 
-def test_delete_when_disconnected(postgres_cache_config_fixture, mocker):
+def test_delete_when_disconnected(postgres_cache_config_fixture, mocker: MockerFixture):
     """Test the delete() method."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -258,7 +277,9 @@ def test_delete_when_disconnected(postgres_cache_config_fixture, mocker):
         cache.delete(USER_ID_1, CONVERSATION_ID_1, False)
 
 
-def test_delete_operation_when_connected(postgres_cache_config_fixture, mocker):
+def test_delete_operation_when_connected(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test the delete() method."""
     # prevent real connection to PG instance
     mock_connect = mocker.patch("psycopg2.connect")
@@ -274,7 +295,9 @@ def test_delete_operation_when_connected(postgres_cache_config_fixture, mocker):
     assert cache.delete(USER_ID_1, CONVERSATION_ID_1, False) is False
 
 
-def test_delete_operation_operation_error(postgres_cache_config_fixture, mocker):
+def test_delete_operation_operation_error(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test the delete() method."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -288,7 +311,9 @@ def test_delete_operation_operation_error(postgres_cache_config_fixture, mocker)
         cache.delete(USER_ID_1, CONVERSATION_ID_1, False)
 
 
-def test_list_operation_when_disconnected(postgres_cache_config_fixture, mocker):
+def test_list_operation_when_disconnected(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test the list() method."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -302,7 +327,9 @@ def test_list_operation_when_disconnected(postgres_cache_config_fixture, mocker)
         cache.list(USER_ID_1, False)
 
 
-def test_list_operation_when_connected(postgres_cache_config_fixture, mocker):
+def test_list_operation_when_connected(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test the list() method."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")
@@ -314,7 +341,7 @@ def test_list_operation_when_connected(postgres_cache_config_fixture, mocker):
     assert isinstance(lst, list)
 
 
-def test_topic_summary_operations(postgres_cache_config_fixture, mocker):
+def test_topic_summary_operations(postgres_cache_config_fixture, mocker: MockerFixture):
     """Test topic summary set operations and retrieval via list."""
     # prevent real connection to PG instance
     mock_connect = mocker.patch("psycopg2.connect")
@@ -343,7 +370,9 @@ def test_topic_summary_operations(postgres_cache_config_fixture, mocker):
     assert isinstance(conversations[0], ConversationData)
 
 
-def test_topic_summary_after_conversation_delete(postgres_cache_config_fixture, mocker):
+def test_topic_summary_after_conversation_delete(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test that topic summary is deleted when conversation is deleted."""
     # prevent real connection to PG instance
     mock_connect = mocker.patch("psycopg2.connect")
@@ -364,7 +393,9 @@ def test_topic_summary_after_conversation_delete(postgres_cache_config_fixture,
     assert deleted is True
 
 
-def test_topic_summary_when_disconnected(postgres_cache_config_fixture, mocker):
+def test_topic_summary_when_disconnected(
+    postgres_cache_config_fixture, mocker: MockerFixture
+):
     """Test topic summary operations when cache is disconnected."""
     # prevent real connection to PG instance
     mocker.patch("psycopg2.connect")

tests/unit/metrics/test_utis.py

@@ -1,9 +1,10 @@
 """Unit tests for functions defined in metrics/utils.py"""
 
+from pytest_mock import MockerFixture
 from metrics.utils import setup_model_metrics, update_llm_token_count_from_turn
 
 
-async def test_setup_model_metrics(mocker) -> None:
+async def test_setup_model_metrics(mocker: MockerFixture) -> None:
     """Test the setup_model_metrics function."""
 
     # Mock the LlamaStackAsLibraryClient
@@ -76,7 +77,7 @@ async def test_setup_model_metrics(mocker) -> None:
     )
 
 
-def test_update_llm_token_count_from_turn(mocker) -> None:
+def test_update_llm_token_count_from_turn(mocker: MockerFixture) -> None:
     """Test the update_llm_token_count_from_turn function."""
     mocker.patch("metrics.utils.Tokenizer.get_instance")
     mock_formatter_class = mocker.patch("metrics.utils.ChatFormat")

tests/unit/utils/auth_helpers.py

@@ -1,6 +1,5 @@
 """Helper functions for mocking authorization in tests."""
 
-from typing import Any
 from unittest.mock import AsyncMock, Mock
 from pytest_mock import MockerFixture
 

tests/unit/utils/test_common.py

@@ -1,8 +1,8 @@
 """Test module for utils/common.py."""
 
 from unittest.mock import Mock, AsyncMock
-from pytest_mock import MockerFixture
 from logging import Logger
+from pytest_mock import MockerFixture
 
 import pytest