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, text
from sqlmodel import Session

from domain.tasks.subtasks import SubTaskStatus
from domain.tasks.subtasks import SubTaskFailure
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_reports_an_ordinarily_failing_record_for_redelivery(
    harness: Harness,
) -> None:
    # Arrange
    @task_handler(
        task_source="abri_api",
        source=Source.HUBSPOT_DEAL,
        orchestrator_cm=harness.factory,
    )
    def handler(body: dict[str, Any], context: Any) -> None:
        raise RuntimeError("transient failure")

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

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

    # Assert
    assert result["batchItemFailures"] == [{"itemIdentifier": "msg-1"}]
    subtask_id = result["tasks"][0]["subtask_id"]
    assert harness.subtasks.get(UUID(subtask_id)).status is SubTaskStatus.FAILED


def test_task_handler_does_not_requeue_a_record_failing_non_retriably(
    harness: Harness,
) -> None:
    # Arrange
    @task_handler(
        task_source="abri_api",
        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_attaches_to_a_supplied_task_and_subtask(
    harness: Harness, db_engine: Engine
) -> None:
    """A message carrying task_id + subtask_id (a Modelling Run batch,
    ADR-0055) runs under the distributor's pre-created sub_task — no new Task
    or SubTask rows are created."""
    # arrange — the distributor's pre-created task + batch sub_task
    task, subtask = harness.orchestrator.create_task_with_subtask(
        task_source="app:modelling_run", inputs={"property_ids": [1, 2]}
    )
    received: list[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(task_id)

    event = {
        "Records": [
            {
                "messageId": "m-1",
                "body": (
                    f'{{"task_id": "{task.id}", "subtask_id": "{subtask.id}", '
                    f'"property_ids": [1, 2]}}'
                ),
            }
        ]
    }

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

    # assert — ran under the supplied ids, sub_task completed, nothing created
    assert received == [task.id]
    assert result["tasks"] == [{"task_id": str(task.id), "subtask_id": str(subtask.id)}]
    assert harness.subtasks.get(subtask.id).status.value == "complete"
    with db_engine.connect() as conn:
        assert conn.execute(text("SELECT count(*) FROM tasks")).scalar_one() == 1
        assert conn.execute(text("SELECT count(*) FROM sub_task")).scalar_one() == 1


def test_recorded_failure_fails_the_subtask_without_an_sqs_retry(
    harness: Harness,
) -> None:
    """A SubTaskFailure is a database record, not retry fuel (ADR-0055): the
    sub_task fails with the structured details, but the message is NOT
    reported to SQS as a batch item failure."""
    # arrange
    task, subtask = harness.orchestrator.create_task_with_subtask(
        task_source="app:modelling_run", inputs={"property_ids": [1, 2]}
    )

    @task_handler(
        task_source="modelling_e2e",
        source=Source.PROPERTY,
        orchestrator_cm=harness.factory,
    )
    def handler(body: dict[str, Any], context: Any) -> None:
        raise SubTaskFailure(
            "1 of 2 properties failed",
            details={"succeeded": 1, "failed": [{"property_id": 2, "error": "x"}]},
        )

    event = {
        "Records": [
            {
                "messageId": "m-1",
                "body": (
                    f'{{"task_id": "{task.id}", "subtask_id": "{subtask.id}", '
                    f'"property_ids": [1, 2]}}'
                ),
            }
        ]
    }

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

    # assert — recorded, not retried
    assert result["batchItemFailures"] == []
    failed = harness.subtasks.get(subtask.id)
    assert failed.status.value == "failed"
    assert failed.outputs is not None and failed.outputs["succeeded"] == 1


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