feat: Refactor agent registry usage and add mock fixtures for testing

Author: HyeockJinKim

src/ai/backend/manager/repositories/model_serving/repository.py

@@ -58,8 +58,7 @@
 # Layer-specific decorator for model_serving repository
 repository_decorator = create_layer_aware_repository_decorator(LayerType.MODEL_SERVING)
 
-if TYPE_CHECKING:
-    from ai.backend.manager.registry import AgentRegistry
+from ai.backend.manager.registry import AgentRegistry
 
 
 class ModelServingRepository:
@@ -727,7 +726,7 @@ async def resolve_image_for_endpoint_creation(
     async def modify_endpoint(
         self,
         action: ModifyEndpointAction,
-        agent_registry: "AgentRegistry",
+        agent_registry: AgentRegistry,
         legacy_etcd_config_loader: LegacyEtcdLoader,
         storage_manager: StorageSessionManager,
     ) -> MutationResult:

tests/manager/repositories/model_serving/test_admin_repository.py

@@ -1,4 +1,5 @@
 import uuid
+from unittest.mock import AsyncMock
 
 import pytest
 from sqlalchemy.orm.exc import NoResultFound
@@ -9,6 +10,13 @@
 from .conftest import assert_update_query_executed
 
 
+@pytest.fixture
+def mock_valkey_live():
+    mock = AsyncMock()
+    mock.store_live_data = AsyncMock()
+    return mock
+
+
 @pytest.mark.asyncio
 async def test_get_endpoint_by_id_force_success(
     admin_model_serving_repository,
@@ -150,6 +158,7 @@ async def test_update_route_traffic_force_success(
     sample_endpoint,
     patch_routing_get,
     patch_endpoint_get,
+    mock_valkey_live,
 ):
     """Test admin force update of route traffic ratio."""
     # Arrange
@@ -161,7 +170,7 @@ async def test_update_route_traffic_force_success(
 
     # Act
     result = await admin_model_serving_repository.update_route_traffic_force(
-        route_id, service_id, new_traffic_ratio
+        mock_valkey_live, route_id, service_id, new_traffic_ratio
     )
 
     # Assert

tests/manager/repositories/model_serving/test_create_endpoint_validated.py

@@ -1,3 +1,5 @@
+from unittest.mock import AsyncMock, MagicMock
+
 import pytest
 
 from .conftest import (
@@ -8,15 +10,24 @@
 )
 
 
+@pytest.fixture
+def mock_agent_registry():
+    """Mock agent registry for testing."""
+    mock = MagicMock()
+    mock.create_appproxy_endpoint = AsyncMock(return_value="https://test-endpoint.example.com")
+    return mock
+
+
 @pytest.mark.asyncio
 async def test_create_endpoint_validated_success(
     model_serving_repository,
     setup_writable_session,
     sample_endpoint,
+    mock_agent_registry,
 ):
     """Test successful creation of an endpoint."""
     # Act
-    result = await model_serving_repository.create_endpoint_validated(sample_endpoint)
+    result = await model_serving_repository.create_endpoint_validated(sample_endpoint, mock_agent_registry)
 
     # Assert
     assert_basic_endpoint_result(result, sample_endpoint)
@@ -48,6 +59,7 @@ async def test_create_endpoint_validated_with_configurations(
     sample_vfolder,
     endpoint_config,
     expected_attrs,
+    mock_agent_registry,
 ):
     """Test creation of endpoints with different configurations."""
     # Arrange
@@ -57,7 +69,7 @@ async def test_create_endpoint_validated_with_configurations(
         endpoint_row = create_full_featured_endpoint(sample_user, sample_image, sample_vfolder)
 
     # Act
-    result = await model_serving_repository.create_endpoint_validated(endpoint_row)
+    result = await model_serving_repository.create_endpoint_validated(endpoint_row, mock_agent_registry)
 
     # Assert
     assert_basic_endpoint_result(result, endpoint_row)
@@ -75,13 +87,14 @@ async def test_create_endpoint_validated_transaction_handling(
     mock_db_engine,
     mock_session,
     sample_endpoint,
+    mock_agent_registry,
 ):
     """Test that creation properly handles database transactions."""
     # Arrange
     setup_db_session_mock(mock_db_engine, mock_session)
 
     # Act
-    result = await model_serving_repository.create_endpoint_validated(sample_endpoint)
+    result = await model_serving_repository.create_endpoint_validated(sample_endpoint, mock_agent_registry)
 
     # Assert
     assert_basic_endpoint_result(result, sample_endpoint)

tests/manager/repositories/model_serving/test_get_endpoint_by_id_validated.py

@@ -1,4 +1,5 @@
 import uuid
+from unittest.mock import AsyncMock
 
 import pytest
 from sqlalchemy.orm.exc import NoResultFound
@@ -9,6 +10,13 @@
 from .conftest import assert_update_query_executed
 
 
+@pytest.fixture
+def mock_valkey_live():
+    mock = AsyncMock()
+    mock.store_live_data = AsyncMock()
+    return mock
+
+
 @pytest.mark.asyncio
 async def test_get_endpoint_by_id_force_success(
     admin_model_serving_repository,
@@ -150,6 +158,7 @@ async def test_update_route_traffic_force_success(
     sample_endpoint,
     patch_routing_get,
     patch_endpoint_get,
+    mock_valkey_live,
 ):
     """Test admin force update of route traffic ratio."""
     # Arrange
@@ -161,7 +170,7 @@ async def test_update_route_traffic_force_success(
 
     # Act
     result = await admin_model_serving_repository.update_route_traffic_force(
-        route_id, service_id, new_traffic_ratio
+        mock_valkey_live, route_id, service_id, new_traffic_ratio
     )
 
     # Assert

tests/manager/services/model_serving/actions/test_update_route.py

@@ -4,7 +4,7 @@
 import aiohttp
 import pytest
 
-from ai.backend.manager.errors.service import RouteNotFound
+from ai.backend.manager.errors.service import ModelServiceNotFound
 from ai.backend.manager.models.user import UserRole
 from ai.backend.manager.services.model_serving.actions.update_route import (
     UpdateRouteAction,
@@ -60,10 +60,10 @@ def mock_get_endpoint_for_appproxy_update(mocker, mock_repositories):
 
 
 @pytest.fixture
-def mock_update_appproxy_endpoint_routes(mocker, mock_repositories):
+def mock_notify_endpoint_route_update_to_appproxy(mocker, mock_agent_registry):
     mock = mocker.patch.object(
-        mock_repositories.repository,
-        "update_appproxy_endpoint_routes",
+        mock_agent_registry,
+        "notify_endpoint_route_update_to_appproxy",
         new_callable=AsyncMock,
     )
     mock.return_value = None
@@ -130,7 +130,7 @@ async def test_update_route(
         mock_update_route_traffic_force,
         mock_update_route_traffic_validated,
         mock_get_endpoint_for_appproxy_update,
-        mock_update_appproxy_endpoint_routes,
+        mock_notify_endpoint_route_update_to_appproxy,
     ):
         # Mock endpoint data for route update
         mock_endpoint_data = MagicMock(
@@ -182,7 +182,7 @@ async def update_route(action: UpdateRouteAction):
                     route_id=uuid.UUID("99999999-9999-9999-9999-999999999999"),
                     traffic_ratio=0.5,
                 ),
-                RouteNotFound,
+                ModelServiceNotFound,
             ),
         ],
     )
@@ -209,7 +209,7 @@ async def test_update_route_appproxy_failure(
         mock_check_requester_access_update_route,
         mock_update_route_traffic_validated,
         mock_get_endpoint_for_appproxy_update,
-        mock_update_appproxy_endpoint_routes,
+        mock_notify_endpoint_route_update_to_appproxy,
     ):
         action = UpdateRouteAction(
             requester_ctx=RequesterCtx(
@@ -228,9 +228,8 @@ async def test_update_route_appproxy_failure(
         mock_get_endpoint_for_appproxy_update.return_value = MagicMock(id=action.service_id)
 
         # Mock AppProxy communication failure
-        mock_update_appproxy_endpoint_routes.side_effect = aiohttp.ClientError("Connection failed")
-
-        result = await model_serving_processors.update_route.wait_for_complete(action)
+        mock_notify_endpoint_route_update_to_appproxy.side_effect = aiohttp.ClientError("Connection failed")
 
-        # Should still return success despite AppProxy failure
-        assert result.success is True
+        # AppProxy failure should propagate as exception
+        with pytest.raises(aiohttp.ClientError, match="Connection failed"):
+            await model_serving_processors.update_route.wait_for_complete(action)

tests/manager/services/model_serving/conftest.py

@@ -87,6 +87,15 @@ def mock_background_task_manager():
     return mock_background_task_manager
 
 
+@pytest.fixture
+def mock_valkey_live():
+    mock = MagicMock()
+    mock.store_live_data = AsyncMock()
+    mock.get_live_data = AsyncMock()
+    mock.delete_live_data = AsyncMock()
+    return mock
+
+
 @pytest.fixture
 def model_serving_service(
     database_fixture,
@@ -96,6 +105,7 @@ def model_serving_service(
     mock_agent_registry,
     mock_background_task_manager,
     mock_config_provider,
+    mock_valkey_live,
     mock_repositories,
 ) -> ModelServingService:
     return ModelServingService(
@@ -104,6 +114,7 @@ def model_serving_service(
         event_dispatcher=mock_event_dispatcher,
         storage_manager=mock_storage_manager,
         config_provider=mock_config_provider,
+        valkey_live=mock_valkey_live,
         repository=mock_repositories.repository,
         admin_repository=mock_repositories.admin_repository,
     )

tests/manager/services/model_serving/fixtures.py

@@ -80,6 +80,15 @@ def mock_background_task_manager():
     return mock_background_task_manager
 
 
+@pytest.fixture
+def mock_valkey_live():
+    mock = MagicMock()
+    mock.store_live_data = AsyncMock()
+    mock.get_live_data = AsyncMock()
+    mock.delete_live_data = AsyncMock()
+    return mock
+
+
 @pytest.fixture
 def mock_service(
     database_fixture,  # noqa: ARG001
@@ -89,6 +98,7 @@ def mock_service(
     mock_agent_registry,
     mock_background_task_manager,
     mock_config_provider,
+    mock_valkey_live,
     mock_repositories,
 ) -> ModelServingService:
     return ModelServingService(
@@ -97,6 +107,7 @@ def mock_service(
         event_dispatcher=mock_event_dispatcher,
         storage_manager=mock_storage_manager,
         config_provider=mock_config_provider,
+        valkey_live=mock_valkey_live,
         repository=mock_repositories.repository,
         admin_repository=mock_repositories.admin_repository,
     )