More refactoring

Author: tisnik

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
 

tests/unit/utils/test_endpoints.py

@@ -2,6 +2,7 @@
 
 import os
 import pytest
+from pytest_mock import MockerFixture
 from fastapi import HTTPException
 from pydantic import AnyUrl
 
@@ -245,7 +246,9 @@ def test_get_profile_prompt_with_enabled_query_system_prompt(
 
 
 @pytest.mark.asyncio
-async def test_get_agent_with_conversation_id(prepare_agent_mocks, mocker):
+async def test_get_agent_with_conversation_id(
+    prepare_agent_mocks, mocker: MockerFixture
+):
     """Test get_agent function when agent exists in llama stack."""
     mock_client, mock_agent = prepare_agent_mocks
     conversation_id = "test_conversation_id"
@@ -664,7 +667,9 @@ async def test_get_agent_no_tools_false_preserves_parser(
 
 
 @pytest.mark.asyncio
-async def test_get_temp_agent_basic_functionality(prepare_agent_mocks, mocker):
+async def test_get_temp_agent_basic_functionality(
+    prepare_agent_mocks, mocker: MockerFixture
+):
     """Test get_temp_agent function creates agent with correct parameters."""
     mock_client, mock_agent = prepare_agent_mocks
     mock_agent.create_session.return_value = "temp_session_id"
@@ -703,7 +708,9 @@ async def test_get_temp_agent_basic_functionality(prepare_agent_mocks, mocker):
 
 
 @pytest.mark.asyncio
-async def test_get_temp_agent_returns_valid_ids(prepare_agent_mocks, mocker):
+async def test_get_temp_agent_returns_valid_ids(
+    prepare_agent_mocks, mocker: MockerFixture
+):
     """Test get_temp_agent function returns valid agent_id and session_id."""
     mock_client, mock_agent = prepare_agent_mocks
     mock_agent.agent_id = "generated_agent_id"
@@ -736,7 +743,9 @@ async def test_get_temp_agent_returns_valid_ids(prepare_agent_mocks, mocker):
 
 
 @pytest.mark.asyncio
-async def test_get_temp_agent_no_persistence(prepare_agent_mocks, mocker):
+async def test_get_temp_agent_no_persistence(
+    prepare_agent_mocks, mocker: MockerFixture
+):
     """Test get_temp_agent function creates agent without session persistence."""
     mock_client, mock_agent = prepare_agent_mocks
     mock_agent.create_session.return_value = "temp_session_id"
@@ -818,7 +827,9 @@ def test_get_topic_summary_system_prompt_with_custom_profile():
     assert topic_summary_prompt == prompts.get("topic_summary")
 
 
-def test_get_topic_summary_system_prompt_with_custom_profile_no_topic_summary(mocker):
+def test_get_topic_summary_system_prompt_with_custom_profile_no_topic_summary(
+    mocker: MockerFixture,
+):
     """Test that default topic summary prompt is returned when custom profile has
     no topic_summary prompt.
     """

tests/unit/utils/test_llama_stack_version.py

@@ -2,6 +2,7 @@
 
 import pytest
 from semver import Version
+from pytest_mock import MockerFixture
 
 from llama_stack_client.types import VersionInfo
 
@@ -17,7 +18,7 @@
 
 
 @pytest.mark.asyncio
-async def test_check_llama_stack_version_minimal_supported_version(mocker):
+async def test_check_llama_stack_version_minimal_supported_version(mocker: MockerFixture):
     """Test the check_llama_stack_version function."""
 
     # mock the Llama Stack client
@@ -31,7 +32,7 @@ async def test_check_llama_stack_version_minimal_supported_version(mocker):
 
 
 @pytest.mark.asyncio
-async def test_check_llama_stack_version_maximal_supported_version(mocker):
+async def test_check_llama_stack_version_maximal_supported_version(mocker: MockerFixture):
     """Test the check_llama_stack_version function."""
 
     # mock the Llama Stack client
@@ -45,7 +46,7 @@ async def test_check_llama_stack_version_maximal_supported_version(mocker):
 
 
 @pytest.mark.asyncio
-async def test_check_llama_stack_version_too_small_version(mocker):
+async def test_check_llama_stack_version_too_small_version(mocker: MockerFixture):
     """Test the check_llama_stack_version function."""
 
     # mock the Llama Stack client