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.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 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_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_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