: add integration test

tests/integration/responses/test_extra_body_shields.py

@@ -0,0 +1,83 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+"""
+Test for extra_body parameter support with shields example.
+
+This test demonstrates that parameters marked with ExtraBodyField annotation
+can be passed via extra_body in the client SDK and are received by the
+server-side implementation.
+"""
+
+import pytest
+from llama_stack_client import APIStatusError
+
+
+def test_shields_via_extra_body(compat_client, text_model_id):
+    """Test that shields parameter is received by the server and raises NotImplementedError."""
+
+    # Test with shields as list of strings (shield IDs)
+    with pytest.raises((APIStatusError, NotImplementedError)) as exc_info:
+        compat_client.responses.create(
+            model=text_model_id,
+            input="What is the capital of France?",
+            stream=False,
+            extra_body={
+                "shields": ["test-shield-1", "test-shield-2"]
+            }
+        )
+
+    # Verify the error message indicates shields are not implemented
+    error_message = str(exc_info.value)
+    assert "not yet implemented" in error_message.lower() or "not implemented" in error_message.lower()
+
+
+
+
+def test_response_without_shields_still_works(compat_client, text_model_id):
+    """Test that responses still work without shields parameter (backwards compatibility)."""
+
+    response = compat_client.responses.create(
+        model=text_model_id,
+        input="Hello, world!",
+        stream=False,
+    )
+
+    # Verify response was created successfully
+    assert response.id is not None
+    assert response.output_text is not None
+    assert len(response.output_text) > 0
+
+
+def test_shields_parameter_received_end_to_end(compat_client, text_model_id):
+    """
+    Test that shields parameter passed via extra_body reaches the server implementation.
+
+    This verifies end-to-end that:
+    1. The parameter can be passed via extra_body in the client SDK
+    2. The parameter is properly routed through the API layers
+    3. The server-side implementation receives the parameter (verified by NotImplementedError)
+
+    The NotImplementedError proves the extra_body parameter reached the implementation,
+    as opposed to being rejected earlier due to signature mismatch or validation errors.
+    """
+    # Test with shields parameter via extra_body
+    with pytest.raises((APIStatusError, NotImplementedError)) as exc_info:
+        compat_client.responses.create(
+            model=text_model_id,
+            input="Test message for shields verification",
+            stream=False,
+            extra_body={
+                "shields": ["shield-1", "shield-2"]
+            }
+        )
+
+    # The NotImplementedError proves that:
+    # 1. extra_body.shields was parsed and passed to the API
+    # 2. The server-side implementation received the shields parameter
+    # 3. No signature mismatch or validation errors occurred
+    error_message = str(exc_info.value)
+    assert "not yet implemented" in error_message.lower() or "not implemented" in error_message.lower()