feat: add stepfunctions contract tests

Author: yiyuan-he

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

@@ -12,6 +12,7 @@
 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
@@ -47,6 +48,8 @@ def do_GET(self):
             self._handle_bedrock_request()
         if self.in_path("secretsmanager"):
             self._handle_secretsmanager_request()
+        if self.in_path("stepfunctions"):
+            self._handle_stepfunctions_request()
 
         self._end_request(self.main_status)
 
@@ -329,6 +332,37 @@ def _handle_secretsmanager_request(self) -> None:
         else:
             set_main_status(404)
 
+    def _handle_stepfunctions_request(self) -> None:
+        sfn_client = boto3.client("stepfunctions", endpoint_url=_AWS_SDK_ENDPOINT, region_name=_AWS_REGION)
+        if self.in_path(_ERROR):
+            set_main_status(400)
+            try:
+                error_client = boto3.client("stepfunctions", endpoint_url=_ERROR_ENDPOINT, region_name=_AWS_REGION)
+                error_client.describe_state_machine(stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:unExistStateMachine")
+            except Exception as exception:
+                print("Expected exception occurred", exception)
+        elif self.in_path(_FAULT):
+            set_main_status(500)
+            try:
+                fault_client = boto3.client("stepfunctions", endpoint_url=_FAULT_ENDPOINT, region_name=_AWS_REGION)
+                fault_client.meta.events.register(
+                    "before-call.stepfunctions.ListStateMachineVersions",
+                    lambda **kwargs: inject_500_error("ListStateMachineVersions", **kwargs),
+                )
+                fault_client.list_state_machine_versions(
+                    stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:invalid-state-machine",
+                )
+            except Exception as exception:
+                print("Expected exception occurred", exception)
+        elif self.in_path("describestatemachine/my-state-machine"):
+            set_main_status(200)
+            sfn_client.describe_state_machine(stateMachineArn="arn:aws:states:us-west-2:000000000000:stateMachine:testStateMachine")
+        elif self.in_path("describeactivity/my-activity"):
+            set_main_status(200)
+            sfn_client.describe_activity(activityArn="arn:aws:states:us-west-2:000000000000:activity:testActivity")
+        else:
+            set_main_status(404)
+
     def _end_request(self, status_code: int):
         self.send_response_only(status_code)
         self.end_headers()
@@ -383,6 +417,60 @@ def prepare_aws_server() -> None:
                 Name="testSecret", SecretString="secretValue", Description="This is a test secret"
             )
 
+        # Set up Step Functions so tests can access a state machine and activity.
+        sfn_client: BaseClient = boto3.client("stepfunctions", endpoint_url=_AWS_SDK_ENDPOINT, region_name=_AWS_REGION)
+        sfn_response = sfn_client.list_state_machines()
+        state_machine_name = "testStateMachine"
+        activity_name = "testActivity"
+        state_machine = next((st for st in sfn_response["stateMachines"] if st["name"] == state_machine_name), None)
+        if not state_machine:
+            # create state machine needs an iam role so we create it here
+            iam_client: BaseClient = boto3.client("iam", endpoint_url=_AWS_SDK_ENDPOINT, region_name=_AWS_REGION)
+            iam_role_name = "testRole"
+            iam_role_arn = None
+            trust_policy = {
+                "Version": "2012-10-17",
+                "Statement": [
+                    {
+                        "Effect": "Allow",
+                        "Principal": {
+                            "Service": "states.amazonaws.com"
+                        },
+                        "Action": "sts:AssumeRole"
+                    }
+                ]
+            }
+            try:
+                iam_response = iam_client.create_role(
+                    RoleName=iam_role_name, AssumeRolePolicyDocument=json.dumps(trust_policy)
+                )
+                iam_client.attach_role_policy(
+                    RoleName=iam_role_name,
+                    PolicyArn="arn:aws:iam::aws:policy/AWSStepFunctionsFullAccess"
+                )
+                print(f"IAM Role '{iam_role_name}' create successfully.")
+                iam_role_arn = iam_response["Role"]["Arn"]
+                sfn_defintion = {
+                    "Comment": "A simple sequential workflow",
+                    "StartAt": "FirstState",
+                    "States": {
+                        "FirstState": {
+                            "Type": "Pass",
+                            "Result": "Hello, World!",
+                            "End": True
+                        }
+                    }
+                }
+                definition_string = json.dumps(sfn_defintion)
+                sfn_client.create_state_machine(
+                    name=state_machine_name,
+                    definition=definition_string,
+                    roleArn=iam_role_arn
+                )
+                sfn_client.create_activity(name=activity_name)
+            except Exception as exception:
+                print("Something went wrong with Step Functions setup", exception)
+
     except Exception as exception:
         print("Unexpected exception occurred", exception)
 
@@ -412,6 +500,16 @@ def inject_200_success(**kwargs):
     return http_response, response_body
 
 
+def inject_500_error(api_name: str, **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

@@ -36,6 +36,8 @@
 _AWS_BEDROCK_DATA_SOURCE_ID: str = "aws.bedrock.data_source.id"
 _GEN_AI_REQUEST_MODEL: str = "gen_ai.request.model"
 _AWS_SECRET_ARN: str = "aws.secretsmanager.secret.arn"
+_AWS_STATE_MACHINE_ARN: str = "aws.stepfunctions.state_machine.arn"
+_AWS_ACTIVITY_ARN: str = "aws.stepfunctions.activity.arn"
 
 
 # pylint: disable=too-many-public-methods
@@ -75,7 +77,7 @@ def set_up_dependency_container(cls):
         cls._local_stack: LocalStackContainer = (
             LocalStackContainer(image="localstack/localstack:3.5.0")
             .with_name("localstack")
-            .with_services("s3", "sqs", "dynamodb", "kinesis", "secretsmanager")
+            .with_services("s3", "sqs", "dynamodb", "kinesis", "secretsmanager", "iam", "stepfunctions")
             .with_env("DEFAULT_REGION", "us-west-2")
             .with_kwargs(network=NETWORK_NAME, networking_config=local_stack_networking_config)
         )
@@ -590,6 +592,82 @@ def test_secretsmanager_fault(self):
             span_name="Secrets Manager.GetSecretValue",
         )
 
+    def test_stepfunctions_describe_state_machine(self):
+        self.do_test_requests(
+            "stepfunctions/describestatemachine/my-state-machine",
+            "GET",
+            200,
+            0,
+            0,
+            rpc_service="SFN",
+            remote_service="AWS::StepFunctions",
+            remote_operation="DescribeStateMachine",
+            remote_resource_type="AWS::StepFunctions::StateMachine",
+            remote_resource_identifier="testStateMachine",
+            cloudformation_primary_identifier="arn:aws:states:us-west-2:000000000000:stateMachine:testStateMachine",
+            request_specific_attributes={
+                _AWS_STATE_MACHINE_ARN: "arn:aws:states:us-west-2:000000000000:stateMachine:testStateMachine",
+            },
+            span_name="SFN.DescribeStateMachine",
+        )
+
+    def test_stepfunctions_describe_activity(self):
+        self.do_test_requests(
+            "stepfunctions/describeactivity/my-activity",
+            "GET",
+            200,
+            0,
+            0,
+            rpc_service="SFN",
+            remote_service="AWS::StepFunctions",
+            remote_operation="DescribeActivity",
+            remote_resource_type="AWS::StepFunctions::Activity",
+            remote_resource_identifier="testActivity",
+            cloudformation_primary_identifier="arn:aws:states:us-west-2:000000000000:activity:testActivity",
+            request_specific_attributes={
+                _AWS_ACTIVITY_ARN: "arn:aws:states:us-west-2:000000000000:activity:testActivity"
+            },
+            span_name="SFN.DescribeActivity",
+        )
+
+    def test_stepfunctions_error(self):
+        self.do_test_requests(
+            "stepfunctions/error",
+            "GET",
+            400,
+            1,
+            0,
+            rpc_service="SFN",
+            remote_service="AWS::StepFunctions",
+            remote_operation="DescribeStateMachine",
+            remote_resource_type="AWS::StepFunctions::StateMachine",
+            remote_resource_identifier="unExistStateMachine",
+            cloudformation_primary_identifier="arn:aws:states:us-west-2:000000000000:stateMachine:unExistStateMachine",
+            request_specific_attributes={
+                _AWS_STATE_MACHINE_ARN: "arn:aws:states:us-west-2:000000000000:stateMachine:unExistStateMachine",
+            },
+            span_name="SFN.DescribeStateMachine",
+        )
+
+    def test_stepfunctions_fault(self):
+        self.do_test_requests(
+            "stepfunctions/fault",
+            "GET",
+            500,
+            0,
+            1,
+            rpc_service="SFN",
+            remote_service="AWS::StepFunctions",
+            remote_operation="ListStateMachineVersions",
+            remote_resource_type="AWS::StepFunctions::StateMachine",
+            remote_resource_identifier="invalid-state-machine",
+            cloudformation_primary_identifier="arn:aws:states:us-west-2:000000000000:stateMachine:invalid-state-machine",
+            request_specific_attributes={
+                _AWS_STATE_MACHINE_ARN: "arn:aws:states:us-west-2:000000000000:stateMachine:invalid-state-machine",
+            },
+            span_name="SFN.ListStateMachineVersions",
+        )
+
     @override
     def _assert_aws_span_attributes(self, resource_scope_spans: List[ResourceScopeSpan], path: str, **kwargs) -> None:
         target_spans: List[Span] = []