Apply feature parity with Java contract tests.

Author: zzhlogin

contract-tests/images/applications/botocore/botocore_server.py

@@ -12,7 +12,6 @@
 import requests
 from botocore.client import BaseClient
 from botocore.config import Config
-from botocore.exceptions import ClientError
 from typing_extensions import Tuple, override
 
 _PORT: int = 8080
@@ -210,7 +209,7 @@ def _handle_kinesis_request(self) -> None:
 
     def _handle_bedrock_request(self) -> None:
         # Localstack does not support Bedrock related services.
-        # we inject inject_200_success and inject_500_error directly into the API call
+        # we inject inject_200_success directly into the API call
         # to make sure we receive http response with expected status code and attributes.
         bedrock_client: BaseClient = boto3.client("bedrock", endpoint_url=_AWS_SDK_ENDPOINT, region_name=_AWS_REGION)
         bedrock_agent_client: BaseClient = boto3.client(
@@ -222,28 +221,20 @@ def _handle_bedrock_request(self) -> None:
         bedrock_agent_runtime_client: BaseClient = boto3.client(
             "bedrock-agent-runtime", endpoint_url=_AWS_SDK_ENDPOINT, region_name=_AWS_REGION
         )
-        if self.in_path(_ERROR):
-            error_client: BaseClient = boto3.client(
-                "bedrock-agent", endpoint_url=_ERROR_ENDPOINT, region_name=_AWS_REGION
+        if self.in_path("getknowledgebase/get_knowledge_base"):
+            set_main_status(200)
+            bedrock_agent_client.meta.events.register(
+                "before-call.bedrock-agent.GetKnowledgeBase",
+                inject_200_success,
             )
-            set_main_status(400)
-            try:
-                error_client.get_knowledge_base(knowledgeBaseId="invalid-knowledge-base-id")
-            except Exception as exception:
-                print("Expected exception occurred", exception)
-        elif self.in_path(_FAULT):
-            set_main_status(500)
-            try:
-                fault_client: BaseClient = boto3.client(
-                    "bedrock-agent", endpoint_url=_FAULT_ENDPOINT, region_name=_AWS_REGION, config=_NO_RETRY_CONFIG
-                )
-                fault_client.meta.events.register(
-                    "before-call.bedrock-agent.GetDataSource",
-                    lambda **kwargs: inject_500_error("GetDataSource", **kwargs),
-                )
-                fault_client.get_data_source(knowledgeBaseId="TESTKBSEID", dataSourceId="DATASURCID")
-            except Exception as exception:
-                print("Expected exception occurred", exception)
+            bedrock_agent_client.get_knowledge_base(knowledgeBaseId="invalid-knowledge-base-id")
+        elif self.in_path("getdatasource/get_data_source"):
+            set_main_status(200)
+            bedrock_agent_client.meta.events.register(
+                "before-call.bedrock-agent.GetDataSource",
+                inject_200_success,
+            )
+            bedrock_agent_client.get_data_source(knowledgeBaseId="TESTKBSEID", dataSourceId="DATASURCID")
         elif self.in_path("getagent/get-agent"):
             set_main_status(200)
             bedrock_agent_client.meta.events.register(
@@ -272,6 +263,18 @@ def _handle_bedrock_request(self) -> None:
                 sessionId="testSessionId",
                 inputText="Invoke agent sample input text",
             )
+        elif self.in_path("retrieve/retrieve"):
+            set_main_status(200)
+            bedrock_agent_runtime_client.meta.events.register(
+                "before-call.bedrock-agent-runtime.Retrieve",
+                inject_200_success,
+            )
+            bedrock_agent_runtime_client.retrieve(
+                knowledgeBaseId="test-knowledge-base-id",
+                retrievalQuery={
+                    "text": "an example of retrieve query",
+                },
+            )
         elif self.in_path("invokemodel/invoke-model"):
             set_main_status(200)
             bedrock_runtime_client.meta.events.register(
@@ -368,16 +371,6 @@ def inject_200_success(**kwargs):
     return http_response, response_body
 
 
-def inject_500_error(api_name, **kwargs):
-    raise ClientError(
-        {
-            "Error": {"Code": "InternalServerError", "Message": "Internal Server Error"},
-            "ResponseMetadata": {"HTTPStatusCode": 500, "RequestId": "mock-request-id"},
-        },
-        api_name,
-    )
-
-
 def main() -> None:
     prepare_aws_server()
     server_address: Tuple[str, int] = ("0.0.0.0", _PORT)

contract-tests/tests/test/amazon/botocore/botocore_test.py

@@ -31,8 +31,8 @@
 _AWS_KINESIS_STREAM_NAME: str = "aws.kinesis.stream.name"
 _AWS_BEDROCK_AGENT_ID: str = "aws.bedrock.agent.id"
 _AWS_BEDROCK_GUARDRAIL_ID: str = "aws.bedrock.guardrail.id"
-_AWS_BEDROCK_KNOWLEDGEBASE_ID: str = "aws.bedrock.knowledge_base.id"
-_AWS_BEDROCK_DATASOURCE_ID: str = "aws.bedrock.data_source.id"
+_AWS_BEDROCK_KNOWLEDGE_BASE_ID: str = "aws.bedrock.knowledge_base.id"
+_AWS_BEDROCK_DATA_SOURCE_ID: str = "aws.bedrock.data_source.id"
 _GEN_AI_REQUEST_MODEL: str = "gen_ai.request.model"
 
 
@@ -377,7 +377,7 @@ def test_kinesis_fault(self):
             span_name="Kinesis.PutRecord",
         )
 
-    def test_bedrock_invoke_model(self):
+    def test_bedrock_runtime_invoke_model(self):
         self.do_test_requests(
             "bedrock/invokemodel/invoke-model",
             "GET",
@@ -395,24 +395,6 @@ def test_bedrock_invoke_model(self):
             span_name="Bedrock Runtime.InvokeModel",
         )
 
-    def test_bedrock_get_agent(self):
-        self.do_test_requests(
-            "bedrock/getagent/get-agent",
-            "GET",
-            200,
-            0,
-            0,
-            rpc_service="Bedrock Agent",
-            remote_service="AWS::Bedrock",
-            remote_operation="GetAgent",
-            remote_resource_type="AWS::Bedrock::Agent",
-            remote_resource_identifier="TESTAGENTID",
-            request_specific_attributes={
-                _AWS_BEDROCK_AGENT_ID: "TESTAGENTID",
-            },
-            span_name="Bedrock Agent.GetAgent",
-        )
-
     def test_bedrock_get_guardrail(self):
         self.do_test_requests(
             "bedrock/getguardrail/get-guardrail",
@@ -431,7 +413,7 @@ def test_bedrock_get_guardrail(self):
             span_name="Bedrock.GetGuardrail",
         )
 
-    def test_bedrock_invoke_agent(self):
+    def test_bedrock_agent_runtime_invoke_agent(self):
         self.do_test_requests(
             "bedrock/invokeagent/invoke_agent",
             "GET",
@@ -449,38 +431,74 @@ def test_bedrock_invoke_agent(self):
             span_name="Bedrock Agent Runtime.InvokeAgent",
         )
 
-    def test_bedrock_error(self):
+    def test_bedrock_agent_runtime_retrieve(self):
         self.do_test_requests(
-            "bedrock/error",
+            "bedrock/retrieve/retrieve",
             "GET",
-            400,
-            1,
+            200,
+            0,
+            0,
+            rpc_service="Bedrock Agent Runtime",
+            remote_service="AWS::Bedrock",
+            remote_operation="Retrieve",
+            remote_resource_type="AWS::Bedrock::KnowledgeBase",
+            remote_resource_identifier="test-knowledge-base-id",
+            request_specific_attributes={
+                _AWS_BEDROCK_KNOWLEDGE_BASE_ID: "test-knowledge-base-id",
+            },
+            span_name="Bedrock Agent Runtime.Retrieve",
+        )
+
+    def test_bedrock_agent_get_agent(self):
+        self.do_test_requests(
+            "bedrock/getagent/get-agent",
+            "GET",
+            200,
+            0,
+            0,
+            rpc_service="Bedrock Agent",
+            remote_service="AWS::Bedrock",
+            remote_operation="GetAgent",
+            remote_resource_type="AWS::Bedrock::Agent",
+            remote_resource_identifier="TESTAGENTID",
+            request_specific_attributes={
+                _AWS_BEDROCK_AGENT_ID: "TESTAGENTID",
+            },
+            span_name="Bedrock Agent.GetAgent",
+        )
+
+    def test_bedrock_agent_get_knowledge_base(self):
+        self.do_test_requests(
+            "bedrock/getknowledgebase/get_knowledge_base",
+            "GET",
+            200,
+            0,
             0,
             rpc_service="Bedrock Agent",
             remote_service="AWS::Bedrock",
             remote_operation="GetKnowledgeBase",
             remote_resource_type="AWS::Bedrock::KnowledgeBase",
             remote_resource_identifier="invalid-knowledge-base-id",
             request_specific_attributes={
-                _AWS_BEDROCK_KNOWLEDGEBASE_ID: "invalid-knowledge-base-id",
+                _AWS_BEDROCK_KNOWLEDGE_BASE_ID: "invalid-knowledge-base-id",
             },
             span_name="Bedrock Agent.GetKnowledgeBase",
         )
 
-    def test_bedrock_fault(self):
+    def test_bedrock_agent_get_data_source(self):
         self.do_test_requests(
-            "bedrock/fault",
+            "bedrock/getdatasource/get_data_source",
             "GET",
-            500,
+            200,
+            0,
             0,
-            1,
             rpc_service="Bedrock Agent",
             remote_service="AWS::Bedrock",
             remote_operation="GetDataSource",
             remote_resource_type="AWS::Bedrock::DataSource",
             remote_resource_identifier="DATASURCID",
             request_specific_attributes={
-                _AWS_BEDROCK_DATASOURCE_ID: "DATASURCID",
+                _AWS_BEDROCK_DATA_SOURCE_ID: "DATASURCID",
             },
             span_name="Bedrock Agent.GetDataSource",
         )