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", 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", 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