mock_aws.py

import boto3
import io
import zipfile
import json
from moto import mock_sqs, mock_sns, mock_lambda, mock_iam
import os


def _process_lambda(func_str):
    zip_output = io.BytesIO()
    zip_file = zipfile.ZipFile(zip_output, 'w', zipfile.ZIP_DEFLATED)
    zip_file.writestr('lambda_function.py', func_str)
    zip_file.close()
    zip_output.seek(0)
    return zip_output.read()


def echo_lambda_zip():
    pfunc = """
import boto3
import json

def lambda_handler(event, context):
    print('handling ', event)

    msg = json.dumps({"foo":"bar"})
    sns = boto3.client('sns', region_name='us-east-1')
    sns.publish(TargetArn='arn:aws:sqs:us-east-1:123456789012:queue1', Message=msg)
    return msg
"""
    return _process_lambda(pfunc)


@mock_sqs
@mock_sns
@mock_lambda
def run():
    sns_client = boto3.client("sns", region_name="us-east-1")
    sqs_client = boto3.client('sqs', region_name='us-east-1')
    sqs_obj = boto3.resource('sqs', region_name='us-east-1')
    lambda_client = boto3.client('lambda', region_name='us-east-1')

    sqs_queue_name = 'queue1'
    topic_name = 'respond-lambda'

    # lambda_client = boto3.client('lambda', 'us-west-2')
    respond_arn = sns_client.create_topic(Name=topic_name).get('TopicArn')

    sqs_client.create_queue(QueueName=sqs_queue_name)
    sqs_queue = sqs_obj.get_queue_by_name(QueueName=sqs_queue_name)
    queue_url = sqs_client.get_queue_url(QueueName=sqs_queue_name)['QueueUrl']
    sqs_queue_attrs = sqs_client.get_queue_attributes(QueueUrl=queue_url,
                                                      AttributeNames=['All'])['Attributes']
    sqs_queue_arn = sqs_queue_attrs['QueueArn']

    sns_client.subscribe(
        TopicArn=respond_arn,
        Protocol='sqs',
        Endpoint=sqs_queue_arn
    )
    print(respond_arn)

    lambda_echo = lambda_client.create_function(
        FunctionName='testFunction',
        Runtime='python3.6',
        Role='arn:iam:test-iam-role:us-west-2:123456789012:test-iam-role',
        Handler='lambda_function.lambda_handler',
        Code={
            'ZipFile': echo_lambda_zip(),
        },
        Description='test lambda function',
        Timeout=5,
        MemorySize=128,
        Publish=True,
    )
    print('invoking lambda')
    lambda_client.invoke(FunctionName='testFunction',
                         InvocationType="Event", Payload=json.dumps({}))

    msgs = sqs_client.receive_message(
        QueueUrl=queue_url,
        MaxNumberOfMessages=1,
        MessageAttributeNames=[
            'All'
        ],
        VisibilityTimeout=0,
        WaitTimeSeconds=3)
    print("sqs msgs: ", msgs)


if __name__ == "__main__":
    print('running')
    run()