From 9f2ffe53aa384f5745f56c4867e7b8d15f4cc79c Mon Sep 17 00:00:00 2001 From: Jazzcort Date: Fri, 24 Jul 2026 16:08:49 -0400 Subject: [PATCH] Wire run_shield_moderation_v2 into Responses API endpoint Add run_shield_moderation_v2 and build_shield to utils/shields.py to run shield moderation through AbstractSafetyCapability instances instead of the Llama Stack client. Update the responses endpoint to call the new function with shield configs directly. Include unit tests covering pass, block, filtering, and error handling paths. --- src/app/endpoints/a2a.py | 2 +- src/app/endpoints/responses.py | 10 +- src/utils/agents/error_handler.py | 103 +++++++++++ src/utils/agents/query.py | 81 +-------- src/utils/agents/streaming.py | 2 +- src/utils/shields.py | 86 ++++++++- .../endpoints/test_responses_integration.py | 2 +- tests/unit/app/endpoints/test_responses.py | 2 +- tests/unit/utils/agents/test_query.py | 2 +- tests/unit/utils/agents/test_streaming.py | 4 +- tests/unit/utils/test_shields.py | 171 +++++++++++++++++- 11 files changed, 355 insertions(+), 110 deletions(-) create mode 100644 src/utils/agents/error_handler.py diff --git a/src/app/endpoints/a2a.py b/src/app/endpoints/a2a.py index 9ac1b7163..ff8cda209 100644 --- a/src/app/endpoints/a2a.py +++ b/src/app/endpoints/a2a.py @@ -60,7 +60,7 @@ from log import get_logger from models.api.requests import QueryRequest from models.config import Action -from utils.agents.query import map_agent_inference_error +from utils.agents.error_handler import map_agent_inference_error from utils.conversation_compaction import apply_compaction_blocking from utils.mcp_headers import McpHeaders, mcp_headers_dependency from utils.pydantic_ai_helpers import build_agent diff --git a/src/app/endpoints/responses.py b/src/app/endpoints/responses.py index 0b8a6cc65..47035eb9b 100644 --- a/src/app/endpoints/responses.py +++ b/src/app/endpoints/responses.py @@ -105,7 +105,7 @@ select_model_for_responses, ) from utils.rh_identity import get_rh_identity_context -from utils.shields import run_shield_moderation +from utils.shields import run_shield_moderation_v2 from utils.suid import ( normalize_conversation_id, ) @@ -424,11 +424,11 @@ async def responses_endpoint_handler( attachments_text = extract_attachments_text(original_request.input) endpoint_path = ENDPOINT_PATH_RESPONSES - moderation_result = await run_shield_moderation( - client, + + moderation_result = await run_shield_moderation_v2( input_text + "\n\n" + attachments_text, - endpoint_path, - original_request.shield_ids, + configuration.configuration.shields, + responses_request.shield_ids, ) filter_server_tools = ( diff --git a/src/utils/agents/error_handler.py b/src/utils/agents/error_handler.py new file mode 100644 index 000000000..4d2fa8cbc --- /dev/null +++ b/src/utils/agents/error_handler.py @@ -0,0 +1,103 @@ +"""Error mapping for agent inference failures to structured API error responses.""" + +from typing import TypeAlias + +from ogx_client import APIConnectionError, APIStatusError +from pydantic_ai.exceptions import ( + AgentRunError, + ContentFilterError, + IncompleteToolCall, + ModelAPIError, + ModelHTTPError, + UnexpectedModelBehavior, + UsageLimitExceeded, +) + +from log import get_logger +from models.api.responses.error import ( + AbstractErrorResponse, + InternalServerErrorResponse, + PromptTooLongResponse, + QuotaExceededResponse, + ServiceUnavailableResponse, +) +from utils.query import ( + handle_known_apistatus_errors, + is_context_length_error, +) + +AgentInferenceError: TypeAlias = ( + AgentRunError | APIStatusError | APIConnectionError | RuntimeError +) + +logger = get_logger(__name__) + + +def map_agent_inference_error( + exc: AgentInferenceError, + model_id: str, +) -> AbstractErrorResponse: + """Map agent run failures from pydantic-ai or Llama Stack to an LCS error response. + + Args: + exc: Agent, HTTP status, connection, or context-length runtime error. + model_id: Model identifier in provider/model format. + + Returns: + Structured error response for HTTP or SSE error events. + + Raises: + RuntimeError: Re-raised when ``exc`` is a non-agent ``RuntimeError`` that is + not a recognized context-length failure. + """ + match exc: + case AgentRunError() as agent_exc: + return map_pydantic_agent_run_error(agent_exc, model_id) + case APIStatusError() as status_exc: + return handle_known_apistatus_errors(status_exc, model_id) + case APIConnectionError() as connection_exc: + return ServiceUnavailableResponse( + backend_name="OGX", + cause=str(connection_exc), + ) + case RuntimeError() as runtime_exc if is_context_length_error(str(runtime_exc)): + return PromptTooLongResponse(model=model_id) + case _: + return InternalServerErrorResponse.generic() + + +def map_pydantic_agent_run_error( + exc: AgentRunError, model_id: str +) -> AbstractErrorResponse: + """Map pydantic-ai ``AgentRunError`` subclasses to LCS error responses. + + Args: + exc: Agent exception to map. + model_id: Model identifier in provider/model format. + + Returns: + Structured error response for HTTP or SSE error events. + """ + match exc: + case ContentFilterError() as filter_exc: + return InternalServerErrorResponse.query_failed(str(filter_exc)) + case IncompleteToolCall(): + return PromptTooLongResponse(model=model_id) + case UnexpectedModelBehavior(): + logger.error("Unexpected model behavior: %s", exc, exc_info=True) + return InternalServerErrorResponse.generic() + case UsageLimitExceeded(): + return QuotaExceededResponse.model(model_id) + case ModelHTTPError() as http_exc if is_context_length_error(str(http_exc)): + return PromptTooLongResponse(model=model_id) + case ModelHTTPError(status_code=429): + return QuotaExceededResponse.model(model_id) + case ModelHTTPError(): + return InternalServerErrorResponse.generic() + case ModelAPIError() as api_exc: + return ServiceUnavailableResponse( + backend_name="OGX", + cause=str(api_exc), + ) + case _: + return InternalServerErrorResponse.query_failed(str(exc)) diff --git a/src/utils/agents/query.py b/src/utils/agents/query.py index 16fa60100..b239dabfe 100644 --- a/src/utils/agents/query.py +++ b/src/utils/agents/query.py @@ -9,12 +9,6 @@ from ogx_client import APIConnectionError, APIStatusError, AsyncOgxClient from pydantic_ai.exceptions import ( AgentRunError, - ContentFilterError, - IncompleteToolCall, - ModelAPIError, - ModelHTTPError, - UnexpectedModelBehavior, - UsageLimitExceeded, ) from pydantic_ai.messages import ModelRequest, ModelResponse, ToolReturnPart from pydantic_ai.run import AgentRunResult @@ -27,14 +21,13 @@ AbstractErrorResponse, InternalServerErrorResponse, PromptTooLongResponse, - QuotaExceededResponse, - ServiceUnavailableResponse, ) from models.common.agents import AgentTurnAccumulator from models.common.moderation import ShieldModerationResult from models.common.responses.responses_api_params import ResponsesApiParams from models.common.responses.types import ResponseInput from models.common.turn_summary import TurnSummary +from utils.agents.error_handler import map_agent_inference_error from utils.agents.tool_processor import ( process_function_tool_call, process_function_tool_result, @@ -45,8 +38,6 @@ from utils.pydantic_ai_helpers import build_agent from utils.query import ( extract_provider_and_model_from_model_id, - handle_known_apistatus_errors, - is_context_length_error, ) from utils.responses import extract_vector_store_ids_from_tools from utils.token_counter import TokenCounter @@ -68,76 +59,6 @@ class AgentFinishReason(str, Enum): ERROR = "error" -def map_agent_inference_error( - exc: AgentInferenceError, - model_id: str, -) -> AbstractErrorResponse: - """Map agent run failures from pydantic-ai or Llama Stack to an LCS error response. - - Args: - exc: Agent, HTTP status, connection, or context-length runtime error. - model_id: Model identifier in provider/model format. - - Returns: - Structured error response for HTTP or SSE error events. - - Raises: - RuntimeError: Re-raised when ``exc`` is a non-agent ``RuntimeError`` that is - not a recognized context-length failure. - """ - match exc: - case AgentRunError() as agent_exc: - return map_pydantic_agent_run_error(agent_exc, model_id) - case APIStatusError() as status_exc: - return handle_known_apistatus_errors(status_exc, model_id) - case APIConnectionError() as connection_exc: - return ServiceUnavailableResponse( - backend_name="OGX", - cause=str(connection_exc), - ) - case RuntimeError() as runtime_exc if is_context_length_error(str(runtime_exc)): - return PromptTooLongResponse(model=model_id) - case _: - return InternalServerErrorResponse.generic() - - -def map_pydantic_agent_run_error( - exc: AgentRunError, model_id: str -) -> AbstractErrorResponse: - """Map pydantic-ai ``AgentRunError`` subclasses to LCS error responses. - - Args: - exc: Agent exception to map. - model_id: Model identifier in provider/model format. - - Returns: - Structured error response for HTTP or SSE error events. - """ - match exc: - case ContentFilterError() as filter_exc: - return InternalServerErrorResponse.query_failed(str(filter_exc)) - case IncompleteToolCall(): - return PromptTooLongResponse(model=model_id) - case UnexpectedModelBehavior(): - logger.error("Unexpected model behavior: %s", exc, exc_info=True) - return InternalServerErrorResponse.generic() - case UsageLimitExceeded(): - return QuotaExceededResponse.model(model_id) - case ModelHTTPError() as http_exc if is_context_length_error(str(http_exc)): - return PromptTooLongResponse(model=model_id) - case ModelHTTPError(status_code=429): - return QuotaExceededResponse.model(model_id) - case ModelHTTPError(): - return InternalServerErrorResponse.generic() - case ModelAPIError() as api_exc: - return ServiceUnavailableResponse( - backend_name="OGX", - cause=str(api_exc), - ) - case _: - return InternalServerErrorResponse.query_failed(str(exc)) - - def get_agent_finish_reason(response: ModelResponse) -> AgentFinishReason: """Get the finish reason from a completed agent model response. diff --git a/src/utils/agents/streaming.py b/src/utils/agents/streaming.py index 63c2d0375..7c36e54a1 100644 --- a/src/utils/agents/streaming.py +++ b/src/utils/agents/streaming.py @@ -43,12 +43,12 @@ from models.common.responses.contexts import ResponseGeneratorContext from models.common.responses.responses_api_params import ResponsesApiParams from models.common.turn_summary import TurnSummary +from utils.agents.error_handler import map_agent_inference_error from utils.agents.query import ( AgentFinishReason, extract_agent_token_usage, get_agent_finish_reason, get_finish_reason_error, - map_agent_inference_error, ) from utils.agents.tool_processor import ( process_function_tool_call, diff --git a/src/utils/shields.py b/src/utils/shields.py index a4470b4dd..4d16dcc20 100644 --- a/src/utils/shields.py +++ b/src/utils/shields.py @@ -3,15 +3,11 @@ from typing import Optional from fastapi import HTTPException -from ogx_client import ( - APIConnectionError, - AsyncOgxClient, -) -from ogx_client import ( - APIStatusError as LLSApiStatusError, -) +from ogx_client import APIConnectionError, APIStatusError, AsyncOgxClient +from pydantic_ai.exceptions import AgentRunError from configuration import AppConfig +from log import get_logger from models.api.requests import QueryRequest from models.api.responses.error import ( InternalServerErrorResponse, @@ -19,8 +15,21 @@ ServiceUnavailableResponse, UnprocessableEntityResponse, ) -from models.common import ShieldModerationPassed, ShieldModerationResult -from models.config import ShieldConfiguration +from models.common.moderation import ( + ShieldModerationPassed, + ShieldModerationResult, +) +from models.config import QuestionValidityConfig, RedactionConfig, ShieldConfiguration +from pydantic_ai_lightspeed.capabilities.base import AbstractSafetyCapability +from pydantic_ai_lightspeed.capabilities.question_validity._capability import ( + QuestionValidity, +) +from pydantic_ai_lightspeed.capabilities.redaction._capability import ( + PiiRedactionCapability, +) +from utils.agents.error_handler import map_agent_inference_error + +logger = get_logger(__name__) def validate_shield_ids_override( @@ -59,6 +68,63 @@ def validate_shield_ids_override( raise HTTPException(**response.model_dump()) +async def run_shield_moderation_v2( + input_text: str, + shield_configs: list[ShieldConfiguration], + selected_shield_ids: Optional[list[str]] = None, +) -> ShieldModerationResult: + """Run v2 shield moderation on input text. + + Iterates through configured shields and runs moderation checks. + + Parameters: + input_text: The text to moderate. + shield_configs: List of shield configurations to evaluate. + selected_shield_ids: Optional list of shield names to filter by. + + Returns: + Result indicating if content was blocked or passed. + """ + selected_shield_configs = get_shields_for_request( + shield_configs, selected_shield_ids + ) + + for shield_config in selected_shield_configs: + shield = build_shield(shield_config) + + try: + shield_result = await shield.run(input_text) + # APIConnectionError and APIStatusError from ogx should not be raised from model_request, + # because they will be caught inside AsyncOpenAI and transferred into openai's + # APIConnectionError. The openai's exceptions will further transferred into ModelHTTPError + # or ModelAPIError by _map_api_errors in OpenAIResponseModel. + except (AgentRunError, RuntimeError) as exc: + model_id = getattr(shield_config.config, "model_id", "unknown-shield-model") + response = map_agent_inference_error(exc, model_id) + raise HTTPException(**response.model_dump()) from exc + + if shield_result.decision == "blocked": + return shield_result + + return ShieldModerationPassed() + + +def build_shield(shield_config: ShieldConfiguration) -> AbstractSafetyCapability: + """Build a safety capability instance from a shield configuration. + + Parameters: + shield_config: The shield configuration to build from. + + Returns: + The constructed safety capability. + """ + match shield_config.config: + case QuestionValidityConfig(): + return QuestionValidity(shield_config.config) + case RedactionConfig(): + return PiiRedactionCapability(shield_config.config) + + async def run_shield_moderation( _client: AsyncOgxClient, _input_text: str, @@ -124,7 +190,7 @@ async def append_turn_to_conversation( cause=str(e), ) raise HTTPException(**error_response.model_dump()) from e - except LLSApiStatusError as e: + except APIStatusError as e: error_response = InternalServerErrorResponse.generic() raise HTTPException(**error_response.model_dump()) from e diff --git a/tests/integration/endpoints/test_responses_integration.py b/tests/integration/endpoints/test_responses_integration.py index aae7b9f26..9ef34711f 100644 --- a/tests/integration/endpoints/test_responses_integration.py +++ b/tests/integration/endpoints/test_responses_integration.py @@ -173,7 +173,7 @@ def _configure_shield_blocked( moderation_id=moderation_id, ) mocker.patch( - "app.endpoints.responses.run_shield_moderation", + "app.endpoints.responses.run_shield_moderation_v2", return_value=blocked, ) diff --git a/tests/unit/app/endpoints/test_responses.py b/tests/unit/app/endpoints/test_responses.py index 57020665b..63b4c203f 100644 --- a/tests/unit/app/endpoints/test_responses.py +++ b/tests/unit/app/endpoints/test_responses.py @@ -193,7 +193,7 @@ def _patch_moderation(mocker: MockerFixture, decision: str = "passed") -> Any: else: moderation_result = ShieldModerationPassed() mocker.patch( - f"{MODULE}.run_shield_moderation", + f"{MODULE}.run_shield_moderation_v2", new=mocker.AsyncMock(return_value=moderation_result), ) return moderation_result diff --git a/tests/unit/utils/agents/test_query.py b/tests/unit/utils/agents/test_query.py index 94f72cf48..f98f4c65a 100644 --- a/tests/unit/utils/agents/test_query.py +++ b/tests/unit/utils/agents/test_query.py @@ -477,7 +477,7 @@ async def test_api_status_error_raises_http_exception( "detail": {"response": "Quota exceeded", "cause": "quota exceeded"}, } mocker.patch( - "utils.agents.query.handle_known_apistatus_errors", + "utils.agents.error_handler.handle_known_apistatus_errors", return_value=mock_error, ) diff --git a/tests/unit/utils/agents/test_streaming.py b/tests/unit/utils/agents/test_streaming.py index 3190feb26..132e4ad13 100644 --- a/tests/unit/utils/agents/test_streaming.py +++ b/tests/unit/utils/agents/test_streaming.py @@ -605,7 +605,7 @@ async def test_agent_error_raises_http_exception( "detail": {"response": "Error", "cause": "agent failed"}, } mocker.patch( - "utils.agents.streaming.map_agent_inference_error", + "utils.agents.error_handler.map_agent_inference_error", return_value=mock_error, ) @@ -742,7 +742,7 @@ async def inner() -> AsyncIterator[str]: mock_error.detail.response = "Quota exceeded" mock_error.detail.cause = "quota exceeded" mocker.patch( - "utils.agents.streaming.map_agent_inference_error", + "utils.agents.error_handler.map_agent_inference_error", return_value=mock_error, ) mocker.patch( diff --git a/tests/unit/utils/test_shields.py b/tests/unit/utils/test_shields.py index 7183e3e89..a1c22a8aa 100644 --- a/tests/unit/utils/test_shields.py +++ b/tests/unit/utils/test_shields.py @@ -2,17 +2,24 @@ import pytest from fastapi import HTTPException, status +from pydantic_ai.exceptions import ModelAPIError, ModelHTTPError from pytest_mock import MockerFixture -from models.config import QuestionValidityConfig, QuestionValidityShieldConfiguration +from models.common.moderation import ShieldModerationBlocked, ShieldModerationPassed +from models.config import ( + QuestionValidityConfig, + QuestionValidityShieldConfiguration, + ShieldConfiguration, +) from utils.shields import ( append_turn_to_conversation, get_shields_for_request, + run_shield_moderation_v2, validate_shield_ids_override, ) -def _shield(name: str) -> QuestionValidityShieldConfiguration: +def _shield_config(name: str) -> QuestionValidityShieldConfiguration: """Build a minimal question-validity shield configuration for tests.""" return QuestionValidityShieldConfiguration( name=name, @@ -133,12 +140,156 @@ def test_raises_422_when_empty_list_shield_ids_and_override_disabled( assert exc_info.value.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY +class TestRunShieldModerationV2: + """Tests for run_shield_moderation_v2 function.""" + + @pytest.mark.asyncio + async def test_returns_passed_when_no_shields(self) -> None: + """Return ShieldModerationPassed when shield list is empty.""" + result = await run_shield_moderation_v2("test input", []) + assert isinstance(result, ShieldModerationPassed) + + @pytest.mark.asyncio + async def test_returns_passed_when_all_shields_pass( + self, mocker: MockerFixture + ) -> None: + """Return ShieldModerationPassed when every shield passes.""" + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock(return_value=ShieldModerationPassed()) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + shields: list[ShieldConfiguration] = [ + _shield_config("s1"), + _shield_config("s2"), + ] + result = await run_shield_moderation_v2("test input", shields) + + assert isinstance(result, ShieldModerationPassed) + assert mock_shield.run.call_count == 2 + + @pytest.mark.asyncio + async def test_returns_blocked_on_first_block(self, mocker: MockerFixture) -> None: + """Return blocked result from first shield that blocks.""" + blocked = ShieldModerationBlocked(message="rejected", moderation_id="modr-123") + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock(return_value=blocked) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + shields: list[ShieldConfiguration] = [ + _shield_config("s1"), + _shield_config("s2"), + ] + result = await run_shield_moderation_v2("test input", shields) + + assert isinstance(result, ShieldModerationBlocked) + assert result.message == "rejected" + mock_shield.run.assert_called_once() + + @pytest.mark.asyncio + async def test_filters_by_selected_shield_ids(self, mocker: MockerFixture) -> None: + """Only run shields matching the selected IDs.""" + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock(return_value=ShieldModerationPassed()) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + shields: list[ShieldConfiguration] = [ + _shield_config("s1"), + _shield_config("s2"), + _shield_config("s3"), + ] + result = await run_shield_moderation_v2( + "test input", shields, selected_shield_ids=["s2"] + ) + + assert isinstance(result, ShieldModerationPassed) + mock_shield.run.assert_called_once() + + @pytest.mark.asyncio + async def test_shields_stops_on_first_block(self, mocker: MockerFixture) -> None: + """Stop at the first blocking shield.""" + blocked = ShieldModerationBlocked(message="rejected", moderation_id="modr-789") + mock_qv_shield = mocker.Mock() + mock_qv_shield.run = mocker.AsyncMock(return_value=blocked) + + mock_redact_shield = mocker.Mock() + mock_redact_shield.run = mocker.AsyncMock(return_value=ShieldModerationPassed()) + + mocker.patch( + "utils.shields.build_shield", + side_effect=[mock_qv_shield, mock_redact_shield], + ) + + shields: list[ShieldConfiguration] = [ + _shield_config("s-1"), + _shield_config("s-2"), + ] + result = await run_shield_moderation_v2("test input", shields) + + assert isinstance(result, ShieldModerationBlocked) + mock_qv_shield.run.assert_called_once() + mock_redact_shield.run.assert_not_called() + + @pytest.mark.asyncio + async def test_raise_503_on_model_api_error(self, mocker: MockerFixture) -> None: + """Raise HTTP 503 when a shield raises ModelAPIError.""" + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock( + side_effect=ModelAPIError("test", "Incompatible mode") + ) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + with pytest.raises(HTTPException) as exc_info: + await run_shield_moderation_v2("test input", [_shield_config("s1")]) + + assert exc_info.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE + assert "OGX" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_raise_429_when_exceeds_quota(self, mocker: MockerFixture) -> None: + """Raise HTTP 429 when a shield raises ModelHTTPError with status 429.""" + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock( + side_effect=ModelHTTPError(429, "openai/gpt-4o-mini", "Quota exceeded") + ) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + with pytest.raises(HTTPException) as exc_info: + await run_shield_moderation_v2("test input", [_shield_config("s1")]) + + assert exc_info.value.status_code == status.HTTP_429_TOO_MANY_REQUESTS + assert "test-model" in str(exc_info.value.detail) + assert "The model quota has been exceeded" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_raise_413_when_exceeds_context_length( + self, mocker: MockerFixture + ) -> None: + """Raise HTTP 413 when a shield raises ModelHTTPError due to context length exceeded.""" + mock_shield = mocker.Mock() + mock_shield.run = mocker.AsyncMock( + side_effect=ModelHTTPError( + 413, "openai/gpt-4o-mini", "Context length exceeded" + ) + ) + mocker.patch("utils.shields.build_shield", return_value=mock_shield) + + with pytest.raises(HTTPException) as exc_info: + await run_shield_moderation_v2("test input", [_shield_config("s1")]) + + assert exc_info.value.status_code == status.HTTP_413_CONTENT_TOO_LARGE + assert "test-model" in str(exc_info.value.detail) + assert "Prompt is too long" in str(exc_info.value.detail) + + class TestGetShieldsForRequest: """Tests for get_shields_for_request function.""" def test_returns_all_shields_when_shield_ids_none(self) -> None: """Return all configured shields when shield_ids is None.""" - shields = [_shield("shield-1"), _shield("shield-2")] + shields = [ + _shield_config("shield-1"), + _shield_config("shield-2"), + ] result = get_shields_for_request(shields, shield_ids=None) @@ -146,7 +297,10 @@ def test_returns_all_shields_when_shield_ids_none(self) -> None: def test_returns_empty_list_when_shield_ids_empty(self) -> None: """Return no shields when an empty shield_ids list is provided.""" - shields = [_shield("shield-1"), _shield("shield-2")] + shields = [ + _shield_config("shield-1"), + _shield_config("shield-2"), + ] result = get_shields_for_request(shields, shield_ids=[]) @@ -154,9 +308,9 @@ def test_returns_empty_list_when_shield_ids_empty(self) -> None: def test_filters_to_requested_shields_when_all_exist(self) -> None: """Return only shields whose names appear in shield_ids.""" - shield1 = _shield("shield-1") - shield2 = _shield("shield-2") - shield3 = _shield("shield-3") + shield1 = _shield_config("shield-1") + shield2 = _shield_config("shield-2") + shield3 = _shield_config("shield-3") result = get_shields_for_request( [shield1, shield2, shield3], shield_ids=["shield-1", "shield-3"] @@ -168,7 +322,8 @@ def test_raises_404_when_requested_shield_not_configured(self) -> None: """Raise 404 when a requested shield name is not configured.""" with pytest.raises(HTTPException) as exc_info: get_shields_for_request( - [_shield("shield-1")], shield_ids=["shield-1", "missing-shield"] + [_shield_config("shield-1")], + shield_ids=["shield-1", "missing-shield"], ) assert exc_info.value.status_code == status.HTTP_404_NOT_FOUND