test_sqs_cats.py

import json
from typing import List, Generic, TypeVar

from pydantic import Json, Field, BaseModel
from pydantic.generics import GenericModel

MessageType = TypeVar('MessageType')
RecordType = TypeVar('RecordType')


class SQSRecord(GenericModel, Generic[RecordType]):
    """Individual SQS event record."""

    body: Json[RecordType]


class SQSEvent(GenericModel, Generic[MessageType]):
    """
    Message received by Lambda function when triggered by SQS.

    Is an array of individual records.
    """

    records: List[SQSRecord[MessageType]] = Field(alias='Records')


class Cat(BaseModel):
    """Cat description."""

    name: str
    color: str


class SQSCatEvent(SQSEvent[Cat]):
    """Cat descriptions from SQS queue."""


def test_sqs_cats_event():
    raw_event = {
        'Records': [{
            'body': json.dumps({
                'Name': 'Spike',
                'color': 'black',
            }),
        }, {
            'body': json.dumps({
                'name': 'Snow',
                'color': 'white',
            }),
        }],
    }

    event = SQSCatEvent(**raw_event)

    record = event.records[0]
    assert isinstance(record.body, Cat), f'{record.body} is not a Cat instance.'


if __name__ == '__main__':
    test_sqs_cats_event()