from unittest.mock import patch

import boto3
from botocore.stub import Stubber


# create an ssm client
ssm = boto3.client("ssm")
# create an sqs client
sqs = boto3.client("sqs")

# create a Stubber response for ssm
ssm_stubber = Stubber(ssm)
ssm_stubber.add_response(
    "get_parameter",
    {
        "Parameter": {
            "Name": "my_parameter",
            "Type": "String",
            "Value": "my_value",
            "Version": 1,
            "LastModifiedDate": 1234567890.123,
            "ARN": "arn:aws:ssm:us-east-1:123456789012:parameter/my_parameter",
        }
    },
    {"Name": "my_parameter", "WithDecryption": True},
)
ssm_stubber.activate()

# create a Stubber response for sqs
sqs_stubber = Stubber(sqs)
sqs_stubber.add_response(
    "get_queue_url",
    {"QueueUrl": "https://sqs.us-east-1.amazonaws.com/123456789012/my_queue"},
    {"QueueName": "my_queue", "QueueOwnerAWSAccountId": "123456789012"},
)
sqs_stubber.activate()

# patch boto3.client with a side_effect that switches based on the client name
with patch("boto3.client") as mock_client:
    mock_client.side_effect = lambda svc_name: {
        "ssm": ssm_stubber.client,
        "sqs": sqs_stubber.client,
    }.get(svc_name, boto3.client(svc_name))

    # run the code that uses the ssm and sqs clients
    ssm_param = ssm.get_parameter(Name="my_parameter", WithDecryption=True)
    sqs_queue = sqs.get_queue_url(
        QueueName="my_queue", QueueOwnerAWSAccountId="123456789012"
    )

    # assert that the stubbed responses were called
    ssm_stubber.assert_no_pending_responses()
    sqs_stubber.assert_no_pending_responses()

assert ssm_param["Parameter"]["Value"] == "my_value"
assert (
    sqs_queue["QueueUrl"] == "https://sqs.us-east-1.amazonaws.com/123456789012/my_queue"
)