Model/tests/utilities/aws_lambda/test_task_handler.py

from collections.abc import Generator, Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any
from uuid import UUID

import pytest
from sqlalchemy import Engine
from sqlmodel import Session

from domain.tasks.subtasks import SubTaskStatus
from domain.tasks.tasks import Source
from orchestration.task_orchestrator import TaskOrchestrator
from repositories.tasks.subtask_postgres_repository import SubTaskPostgresRepository
from repositories.tasks.task_postgres_repository import TaskPostgresRepository
from utilities.aws_lambda.task_handler import NonRetriableTaskError, task_handler


@dataclass
class Harness:
    orchestrator: TaskOrchestrator
    tasks: TaskPostgresRepository
    subtasks: SubTaskPostgresRepository

    @contextmanager
    def factory(self) -> Generator[TaskOrchestrator, None, None]:
        yield self.orchestrator


@pytest.fixture
def harness(db_engine: Engine) -> Iterator[Harness]:
    with Session(db_engine) as session:
        tasks = TaskPostgresRepository(session=session)
        subtasks = SubTaskPostgresRepository(session=session)
        yield Harness(
            orchestrator=TaskOrchestrator(task_repo=tasks, subtask_repo=subtasks),
            tasks=tasks,
            subtasks=subtasks,
        )


def _direct_event(property_id: str) -> dict[str, Any]:
    return {"property_id": property_id}


def test_task_handler_records_cloudwatch_url_on_subtask(
    harness: Harness, monkeypatch: pytest.MonkeyPatch
) -> None:
    # arrange
    monkeypatch.setenv("AWS_REGION", "eu-west-2")
    monkeypatch.setenv(
        "AWS_LAMBDA_LOG_GROUP_NAME", "/aws/lambda/modelling-e2e"
    )
    monkeypatch.setenv(
        "AWS_LAMBDA_LOG_STREAM_NAME", "2026/05/20/[$LATEST]abc123"
    )

    @task_handler(
        task_source="modelling_e2e",
        source=Source.PROPERTY,
        orchestrator_cm=harness.factory,
    )
    def handler(body: dict[str, Any], context: Any) -> None:
        return None

    # act
    result = handler(_direct_event("prop-1"), context=None)

    # assert
    subtask_id = result[0]["subtask_id"]
    saved_url = harness.subtasks.get(UUID(subtask_id)).cloud_logs_url
    assert saved_url is not None
    assert saved_url.startswith(
        "https://eu-west-2.console.aws.amazon.com/cloudwatch/home"
    )
    # Log group / stream are console-encoded ("/" -> "$252F").
    assert "$252Faws$252Flambda$252Fmodelling-e2e" in saved_url
    assert "$255B$2524LATEST$255D" in saved_url


def test_task_handler_passes_orchestrator_and_task_id_when_flag_is_true(
    harness: Harness,
) -> None:
    # Arrange
    received: list[tuple[TaskOrchestrator, UUID]] = []

    @task_handler(
        task_source="modelling_e2e",
        source=Source.PROPERTY,
        orchestrator_cm=harness.factory,
        pass_task_orchestrator=True,
    )
    def _handler(
        body: dict[str, Any], context: Any, orchestrator: TaskOrchestrator, task_id: UUID
    ) -> None:
        received.append((orchestrator, task_id))

    # Act
    result = _handler(_direct_event("prop-1"), context=None)

    # Assert
    assert len(received) == 1
    recv_orchestrator, recv_task_id = received[0]
    assert recv_orchestrator is harness.orchestrator
    assert recv_task_id == UUID(result[0]["task_id"])


def test_task_handler_does_not_requeue_a_record_failing_non_retriably(
    harness: Harness,
) -> None:
    # Arrange
    @task_handler(
        task_source="abri",
        source=Source.HUBSPOT_DEAL,
        orchestrator_cm=harness.factory,
    )
    def handler(body: dict[str, Any], context: Any) -> None:
        raise NonRetriableTaskError("job logged but write-back failed")

    event = {
        "Records": [
            {"messageId": "msg-1", "body": '{"hubspot_deal_id": "123"}'}
        ]
    }

    # Act
    result = handler(event, context=None)

    # Assert: the failure is recorded, but the message is not requeued.
    assert result["batchItemFailures"] == []
    subtask_id = result["tasks"][0]["subtask_id"]
    failed = harness.subtasks.get(UUID(subtask_id))
    assert failed.status is SubTaskStatus.FAILED
    assert failed.outputs == {"error": "job logged but write-back failed"}


def test_task_handler_leaves_cloudwatch_url_unset_outside_lambda(
    harness: Harness, monkeypatch: pytest.MonkeyPatch
) -> None:
    # arrange
    for var in (
        "AWS_REGION",
        "AWS_LAMBDA_LOG_GROUP_NAME",
        "AWS_LAMBDA_LOG_STREAM_NAME",
    ):
        monkeypatch.delenv(var, raising=False)

    @task_handler(
        task_source="modelling_e2e",
        source=Source.PROPERTY,
        orchestrator_cm=harness.factory,
    )
    def handler(body: dict[str, Any], context: Any) -> None:
        return None

    # act
    result = handler(_direct_event("prop-1"), context=None)

    # assert
    subtask_id = result[0]["subtask_id"]
    assert harness.subtasks.get(UUID(subtask_id)).cloud_logs_url is None