"""Tests for the InferenceGateway facade, lineage recording, and adapters. Covers: - Gateway creates correct client type per protocol - Gateway reuses clients for same endpoint - Target refresh invalidates cached client - Lineage recording captures all required fields - Extraction adapter returns lineage metadata - Unknown protocols fail closed Requirements: 2.1, 2.6, 2.12, 13.6 """ from __future__ import annotations import uuid from unittest.mock import AsyncMock, patch import pytest from services.shared.inference.gateway import InferenceGateway from services.shared.inference.lineage import ( build_lineage_from_result, lineage_to_persistence_dict, ) from services.shared.inference.models import ( ChatMessage, InferenceResult, InferenceTarget, ModelLineage, ProviderCapabilities, StructuredGenerationRequest, TokenUsage, ) # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- def _make_target( protocol: str = "openai_chat", endpoint_id: uuid.UUID | None = None, deployment_id: uuid.UUID | None = None, model: str = "test-model", ) -> InferenceTarget: """Create a minimal InferenceTarget for testing.""" return InferenceTarget( endpoint_id=endpoint_id or uuid.uuid4(), deployment_id=deployment_id or uuid.uuid4(), protocol=protocol, base_url="http://localhost:8000", model=model, capabilities=ProviderCapabilities( chat_completions=True, json_schema=True, usage=True, ), ) def _make_request() -> StructuredGenerationRequest: """Create a minimal StructuredGenerationRequest.""" return StructuredGenerationRequest( messages=[ChatMessage(role="user", content="hello")], max_output_tokens=256, ) def _make_inference_result( endpoint_id: uuid.UUID | None = None, deployment_id: uuid.UUID | None = None, model: str = "test-model", protocol: str = "openai_chat", ) -> InferenceResult: """Create a typical InferenceResult.""" return InferenceResult( content='{"answer": 42}', parsed={"answer": 42}, endpoint_id=endpoint_id or uuid.uuid4(), deployment_id=deployment_id or uuid.uuid4(), model=model, protocol=protocol, structured_mode="json_schema", latency_ms=150, usage=TokenUsage(input_tokens=50, output_tokens=20, total_tokens=70), request_id="req-001", retries=0, ) # --------------------------------------------------------------------------- # Gateway: correct client type per protocol # --------------------------------------------------------------------------- class TestGatewayClientCreation: """Gateway creates the correct client type based on protocol.""" @pytest.mark.asyncio async def test_creates_openai_client_for_openai_chat(self) -> None: """openai_chat protocol creates OpenAICompatibleClient.""" gateway = InferenceGateway() target = _make_target(protocol="openai_chat") with patch( "services.shared.inference.clients.openai_compatible.OpenAICompatibleClient.generate", new_callable=AsyncMock, return_value=_make_inference_result( endpoint_id=target.endpoint_id, deployment_id=target.deployment_id, ), ): result = await gateway.generate(target, _make_request()) assert result.protocol == "openai_chat" await gateway.close() @pytest.mark.asyncio async def test_creates_ollama_client_for_ollama_native(self) -> None: """ollama_native protocol creates OllamaNativeClient.""" gateway = InferenceGateway() target = _make_target(protocol="ollama_native") with patch( "services.shared.inference.clients.ollama_native.OllamaNativeClient.generate", new_callable=AsyncMock, return_value=_make_inference_result( endpoint_id=target.endpoint_id, deployment_id=target.deployment_id, protocol="ollama_native", ), ): result = await gateway.generate(target, _make_request()) assert result.protocol == "ollama_native" await gateway.close() @pytest.mark.asyncio async def test_unknown_protocol_fails_closed(self) -> None: """Unknown protocol raises ValueError — never silently routes to Ollama.""" gateway = InferenceGateway() # Use a type-ignore here since we're intentionally passing an invalid protocol target = InferenceTarget( endpoint_id=uuid.uuid4(), deployment_id=uuid.uuid4(), protocol="unknown_protocol", # type: ignore[arg-type] base_url="http://localhost:8000", model="test", capabilities=ProviderCapabilities(), ) with pytest.raises(ValueError, match="Unknown inference protocol"): await gateway.generate(target, _make_request()) await gateway.close() # --------------------------------------------------------------------------- # Gateway: client reuse # --------------------------------------------------------------------------- class TestGatewayClientReuse: """Gateway reuses clients for the same endpoint.""" @pytest.mark.asyncio async def test_reuses_client_for_same_endpoint(self) -> None: """Repeated calls with the same target reuse the cached client.""" gateway = InferenceGateway() endpoint_id = uuid.uuid4() target = _make_target(endpoint_id=endpoint_id) mock_result = _make_inference_result(endpoint_id=endpoint_id) with patch( "services.shared.inference.clients.openai_compatible.OpenAICompatibleClient.generate", new_callable=AsyncMock, return_value=mock_result, ): await gateway.generate(target, _make_request()) await gateway.generate(target, _make_request()) # Only one client should be cached assert len(gateway._clients) == 1 assert endpoint_id in gateway._clients await gateway.close() @pytest.mark.asyncio async def test_creates_separate_clients_for_different_endpoints(self) -> None: """Different endpoint IDs get separate cached clients.""" gateway = InferenceGateway() target_a = _make_target(endpoint_id=uuid.uuid4()) target_b = _make_target(endpoint_id=uuid.uuid4()) mock_result_a = _make_inference_result(endpoint_id=target_a.endpoint_id) mock_result_b = _make_inference_result(endpoint_id=target_b.endpoint_id) with patch( "services.shared.inference.clients.openai_compatible.OpenAICompatibleClient.generate", new_callable=AsyncMock, side_effect=[mock_result_a, mock_result_b], ): await gateway.generate(target_a, _make_request()) await gateway.generate(target_b, _make_request()) assert len(gateway._clients) == 2 assert target_a.endpoint_id in gateway._clients assert target_b.endpoint_id in gateway._clients await gateway.close() # --------------------------------------------------------------------------- # Gateway: target refresh # --------------------------------------------------------------------------- class TestGatewayTargetRefresh: """Target refresh invalidates cached client.""" @pytest.mark.asyncio async def test_refresh_invalidates_cached_client(self) -> None: """After refresh, the next call creates a new client.""" gateway = InferenceGateway() endpoint_id = uuid.uuid4() target = _make_target(endpoint_id=endpoint_id) mock_result = _make_inference_result(endpoint_id=endpoint_id) with patch( "services.shared.inference.clients.openai_compatible.OpenAICompatibleClient.generate", new_callable=AsyncMock, return_value=mock_result, ), patch( "services.shared.inference.clients.openai_compatible.OpenAICompatibleClient.close", new_callable=AsyncMock, ) as mock_close: await gateway.generate(target, _make_request()) assert endpoint_id in gateway._clients await gateway.refresh_target(endpoint_id) assert endpoint_id not in gateway._clients mock_close.assert_called_once() await gateway.close() @pytest.mark.asyncio async def test_refresh_nonexistent_endpoint_is_safe(self) -> None: """Refreshing an endpoint that isn't cached does nothing.""" gateway = InferenceGateway() # Should not raise await gateway.refresh_target(uuid.uuid4()) await gateway.close() # --------------------------------------------------------------------------- # Lineage recording # --------------------------------------------------------------------------- class TestLineageRecording: """Lineage recording captures all required fields.""" def test_build_lineage_from_result_captures_all_fields(self) -> None: """All lineage fields are extracted from InferenceResult.""" eid = uuid.uuid4() did = uuid.uuid4() result = InferenceResult( content="test", endpoint_id=eid, deployment_id=did, model="qwen-9b", protocol="openai_chat", structured_mode="json_schema", latency_ms=300, request_id="req-abc", retries=1, ) lineage = build_lineage_from_result(result, trace_id="trace-123") assert lineage.endpoint_id == eid assert lineage.deployment_id == did assert lineage.model == "qwen-9b" assert lineage.protocol == "openai_chat" assert lineage.structured_mode == "json_schema" assert lineage.request_id == "req-abc" assert lineage.latency_ms == 300 assert lineage.retries == 1 assert lineage.trace_id == "trace-123" def test_lineage_to_persistence_dict_maps_protocol(self) -> None: """Protocol is mapped to human-friendly model_provider for persistence.""" lineage = ModelLineage( endpoint_id=uuid.uuid4(), deployment_id=uuid.uuid4(), model="test-model", protocol="openai_chat", structured_mode="json_schema", request_id="req-1", latency_ms=100, retries=0, trace_id="t-1", ) d = lineage_to_persistence_dict(lineage) assert d["model_provider"] == "openai_compatible" assert d["model_name"] == "test-model" assert d["protocol"] == "openai_chat" assert d["endpoint_id"] is not None assert d["deployment_id"] is not None assert d["structured_mode"] == "json_schema" assert d["request_id"] == "req-1" assert d["latency_ms"] == 100 assert d["retries"] == 0 assert d["trace_id"] == "t-1" def test_lineage_ollama_protocol_maps_to_ollama_provider(self) -> None: """ollama_native protocol maps to 'ollama' provider.""" lineage = ModelLineage( model="qwen-9b", protocol="ollama_native", ) d = lineage_to_persistence_dict(lineage) assert d["model_provider"] == "ollama" def test_lineage_specialist_protocol_maps_to_specialist_provider(self) -> None: """specialist_http protocol maps to 'specialist' provider.""" lineage = ModelLineage( model="gliner2-large", protocol="specialist_http", ) d = lineage_to_persistence_dict(lineage) assert d["model_provider"] == "specialist" def test_lineage_serialization(self) -> None: """ModelLineage serializes all fields via model_dump().""" lineage = ModelLineage( endpoint_id=uuid.uuid4(), deployment_id=uuid.uuid4(), model="test", protocol="openai_chat", structured_mode="json_object", request_id="r-1", latency_ms=42, retries=2, trace_id="t-abc", ) data = lineage.model_dump() assert data["model"] == "test" assert data["protocol"] == "openai_chat" assert data["trace_id"] == "t-abc" assert data["latency_ms"] == 42 # --------------------------------------------------------------------------- # Extraction adapter: lineage metadata # --------------------------------------------------------------------------- class TestExtractionAdapterLineage: """Extraction adapter returns lineage metadata.""" @pytest.mark.asyncio async def test_extract_document_returns_lineage(self) -> None: """extract_document bundles lineage with the extraction response.""" from services.extractor.inference_adapter import extract_document gateway = InferenceGateway() endpoint_id = uuid.uuid4() deployment_id = uuid.uuid4() target = _make_target( endpoint_id=endpoint_id, deployment_id=deployment_id, ) # Mock the gateway to return a valid extraction JSON extraction_json = '{"summary":"test","companies":[],"macro_themes":[],"novelty_score":0.5,"confidence":0.8,"extraction_warnings":[]}' mock_result = InferenceResult( content=extraction_json, endpoint_id=endpoint_id, deployment_id=deployment_id, model="test-model", protocol="openai_chat", structured_mode="json_schema", latency_ms=200, usage=TokenUsage(input_tokens=100, output_tokens=50), request_id="req-ext-1", retries=0, ) with patch.object( gateway, "generate", new_callable=AsyncMock, return_value=mock_result, ): result = await extract_document( gateway=gateway, target=target, document_text="Test document text for extraction.", document_id="doc-123", ) # Verify lineage is populated assert result.lineage is not None assert result.lineage.endpoint_id == endpoint_id assert result.lineage.deployment_id == deployment_id assert result.lineage.model == "test-model" assert result.lineage.protocol == "openai_chat" assert result.lineage.request_id == "req-ext-1" assert result.lineage.latency_ms == 200 await gateway.close() @pytest.mark.asyncio async def test_extract_document_lineage_on_failure(self) -> None: """Lineage is still captured even when extraction fails.""" from services.extractor.inference_adapter import extract_document gateway = InferenceGateway() endpoint_id = uuid.uuid4() deployment_id = uuid.uuid4() target = _make_target( endpoint_id=endpoint_id, deployment_id=deployment_id, ) # Mock the gateway to return an error result mock_result = InferenceResult( content="", endpoint_id=endpoint_id, deployment_id=deployment_id, model="test-model", protocol="openai_chat", structured_mode="json_schema", latency_ms=5000, request_id="req-timeout", retries=3, error="Request timed out", error_category="timeout", ) with patch.object( gateway, "generate", new_callable=AsyncMock, return_value=mock_result, ): result = await extract_document( gateway=gateway, target=target, document_text="Some text", document_id="doc-fail", max_retries=0, # no retries for test speed ) # Extraction failed but lineage is still captured assert not result.response.success assert result.lineage.endpoint_id == endpoint_id assert result.lineage.model == "test-model" assert result.lineage.protocol == "openai_chat" await gateway.close() # --------------------------------------------------------------------------- # Gateway: active_endpoints property # --------------------------------------------------------------------------- class TestGatewayProperties: """Gateway exposes useful state for monitoring.""" @pytest.mark.asyncio async def test_active_endpoints_tracks_cached_clients(self) -> None: """active_endpoints shows all cached endpoint IDs.""" gateway = InferenceGateway() eid = uuid.uuid4() target = _make_target(endpoint_id=eid) mock_result = _make_inference_result(endpoint_id=eid) with patch( "services.shared.inference.clients.openai_compatible.OpenAICompatibleClient.generate", new_callable=AsyncMock, return_value=mock_result, ): await gateway.generate(target, _make_request()) assert eid in gateway.active_endpoints assert len(gateway.active_endpoints) == 1 await gateway.close() assert len(gateway.active_endpoints) == 0