File: test_lambda.py

package info (click to toggle)
python-aiobotocore 2.25.2-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 2,524 kB
  • sloc: python: 15,437; makefile: 84
file content (81 lines) | stat: -rw-r--r-- 2,228 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
import base64
import io
import json
import zipfile

# Third Party
import botocore.client

try:
    import httpx
except ImportError:
    httpx = None
import pytest


async def _get_role_arn(iam_client, role_name: str):
    try:
        response = await iam_client.get_role(RoleName=role_name)
        return response["Role"]["Arn"]
    except botocore.client.ClientError:
        response = await iam_client.create_role(
            RoleName=role_name,
            AssumeRolePolicyDocument="some policy",
            Path="/my-path/",
        )
        return response["Role"]["Arn"]


def _process_lambda(func_str) -> bytes:
    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()


@pytest.fixture
def aws_lambda_zip() -> bytes:
    lambda_src = """
import json
def lambda_handler(event, context):
    print(event)
    return {"statusCode": 200, "body": event}
"""
    return _process_lambda(lambda_src)


async def test_run_lambda(
    iam_client, lambda_client, aws_lambda_zip, current_http_backend
):
    function_name = f'test-function-{current_http_backend}'
    role_arn = await _get_role_arn(iam_client, 'test-iam-role')
    lambda_response = await lambda_client.create_function(
        FunctionName=function_name,
        Runtime='python3.9',
        Role=role_arn,
        Handler='lambda_function.lambda_handler',
        Timeout=10,
        MemorySize=128,
        Publish=True,
        Code={'ZipFile': aws_lambda_zip},
    )
    assert lambda_response['FunctionName'] == function_name

    invoke_response = await lambda_client.invoke(
        FunctionName=function_name,
        InvocationType="RequestResponse",
        LogType='Tail',
        Payload=json.dumps({"hello": "world"}),
    )

    async with invoke_response['Payload'] as stream:
        data = await stream.read()

    log_result = base64.b64decode(invoke_response["LogResult"])

    assert json.loads(data) == {'statusCode': 200, "body": {"hello": "world"}}
    assert b"{'hello': 'world'}" in log_result

    await lambda_client.delete_function(FunctionName=function_name)