Model/tests/infrastructure/test_address2uprn_queue_client.py

import json
from collections.abc import Iterator
from typing import Any, cast
from uuid import uuid4

import pytest
from moto import mock_aws

from infrastructure.address2uprn_queue_client import Address2UprnQueueClient
from tests.infrastructure import make_boto_client


@pytest.fixture
def queue_setup() -> Iterator[tuple[Address2UprnQueueClient, Any, str]]:
    with mock_aws():
        boto_client = make_boto_client("sqs")
        queue: dict[str, Any] = boto_client.create_queue(
            QueueName="address2uprn-queue"
        )
        queue_url = cast(str, queue["QueueUrl"])
        yield (
            Address2UprnQueueClient(boto_client, queue_url),
            boto_client,
            queue_url,
        )


def test_publish_returns_message_id(
    queue_setup: tuple[Address2UprnQueueClient, Any, str],
) -> None:
    # arrange
    client, _boto, _url = queue_setup
    # act
    message_id = client.publish(
        parent_task_id=uuid4(),
        child_subtask_id=uuid4(),
        s3_uri="s3://my-bucket/path/to/chunk.csv",
    )
    # assert
    assert isinstance(message_id, str)
    assert message_id


def test_publish_body_uses_typed_shape(
    queue_setup: tuple[Address2UprnQueueClient, Any, str],
) -> None:
    # arrange
    client, boto_client, queue_url = queue_setup
    parent_id = uuid4()
    child_id = uuid4()
    s3_uri = "s3://my-bucket/path/to/chunk.csv"

    # act
    client.publish(
        parent_task_id=parent_id,
        child_subtask_id=child_id,
        s3_uri=s3_uri,
    )

    # assert
    received: dict[str, Any] = boto_client.receive_message(
        QueueUrl=queue_url, MaxNumberOfMessages=1
    )
    messages: list[dict[str, Any]] = received["Messages"]
    assert len(messages) == 1
    body = json.loads(messages[0]["Body"])
    assert body == {
        "task_id": str(parent_id),
        "sub_task_id": str(child_id),
        "s3_uri": s3_uri,
    }