Model/tests/orchestration/test_task_orchestrator.py
Khalim Conn-Kowlessar 561b8b078b Test names state the surviving failure-recovery behaviors 🟪
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-07 11:12:12 +00:00

304 lines
10 KiB
Python

from collections.abc import Iterator
from dataclasses import dataclass
import pytest
from sqlalchemy import Engine
from sqlmodel import Session
from domain.tasks.subtasks import SubTask, SubTaskStatus
from domain.tasks.tasks import Source, TaskStatus
from orchestration.task_orchestrator import TaskOrchestrator
from repositories.tasks.subtask_postgres_repository import SubTaskPostgresRepository
from repositories.tasks.task_postgres_repository import TaskPostgresRepository
@dataclass
class Harness:
orchestrator: TaskOrchestrator
tasks: TaskPostgresRepository
subtasks: SubTaskPostgresRepository
@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 test_create_task_with_subtask_creates_both_in_waiting(
harness: Harness,
) -> None:
# act
task, subtask = harness.orchestrator.create_task_with_subtask(
task_source="manual:test",
inputs={"foo": "bar"},
source=Source.PORTFOLIO,
source_id="abc",
)
# assert
assert task.status is TaskStatus.WAITING
assert subtask.status is SubTaskStatus.WAITING
assert subtask.task_id == task.id
assert subtask.inputs == {"foo": "bar"}
def test_start_subtask_cascades_to_in_progress(harness: Harness) -> None:
# arrange
task, subtask = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
# act
started = harness.orchestrator.start_subtask(
subtask.id, cloud_logs_url="https://example/log"
)
# assert
assert started.status is SubTaskStatus.IN_PROGRESS
assert started.cloud_logs_url == "https://example/log"
assert harness.tasks.get(task.id).status is TaskStatus.IN_PROGRESS
def test_complete_subtask_cascades_to_complete(harness: Harness) -> None:
# arrange
task, subtask = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
harness.orchestrator.start_subtask(subtask.id)
# act
harness.orchestrator.complete_subtask(subtask.id, {"value": 42})
# assert
done_subtask = harness.subtasks.get(subtask.id)
done_task = harness.tasks.get(task.id)
assert done_subtask.outputs == {"result": {"value": 42}}
assert done_task.status is TaskStatus.COMPLETE
assert done_task.job_completed is not None
def test_fail_subtask_cascades_to_failed(harness: Harness) -> None:
# arrange
task, subtask = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
# act
harness.orchestrator.fail_subtask(subtask.id, RuntimeError("boom"))
# assert
failed_subtask = harness.subtasks.get(subtask.id)
failed_task = harness.tasks.get(task.id)
assert failed_subtask.outputs == {"error": "boom"}
assert failed_task.status is TaskStatus.FAILED
def test_failed_subtask_locks_task_failed_even_with_others_complete(
harness: Harness,
) -> None:
# arrange
task, first = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
second = SubTask.create(task_id=task.id)
harness.subtasks.create(second)
# act
harness.orchestrator.complete_subtask(first.id)
harness.orchestrator.fail_subtask(second.id, RuntimeError("nope"))
# assert
assert harness.tasks.get(task.id).status is TaskStatus.FAILED
def test_mixed_complete_and_in_progress_keeps_task_in_progress(
harness: Harness,
) -> None:
# arrange
task, first = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
second = SubTask.create(task_id=task.id)
harness.subtasks.create(second)
# act
harness.orchestrator.complete_subtask(first.id)
harness.orchestrator.start_subtask(second.id)
# assert
assert harness.tasks.get(task.id).status is TaskStatus.IN_PROGRESS
def test_run_subtask_happy_path_returns_result_and_cascades_complete(
harness: Harness,
) -> None:
# arrange
task, subtask = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
# act
result = harness.orchestrator.run_subtask(subtask.id, work=lambda: {"answer": 42})
# assert
assert result == {"answer": 42}
assert harness.subtasks.get(subtask.id).status is SubTaskStatus.COMPLETE
assert harness.tasks.get(task.id).status is TaskStatus.COMPLETE
def test_create_child_subtask_adds_waiting_child_without_changing_parent_status(
harness: Harness,
) -> None:
# arrange
task, first = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
harness.orchestrator.start_subtask(first.id)
assert harness.tasks.get(task.id).status is TaskStatus.IN_PROGRESS
# act
child = harness.orchestrator.create_child_subtask(
task.id, inputs={"split": "a"}
)
# assert
persisted_child = harness.subtasks.get(child.id)
assert persisted_child.task_id == task.id
assert persisted_child.status is SubTaskStatus.WAITING
assert persisted_child.inputs == {"split": "a"}
assert persisted_child.id != first.id
# Cascade is a no-op: parent stays IN_PROGRESS.
assert harness.tasks.get(task.id).status is TaskStatus.IN_PROGRESS
def test_run_subtask_failing_work_marks_failed_and_reraises(
harness: Harness,
) -> None:
# arrange
task, subtask = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
def boom() -> None:
raise RuntimeError("boom")
# act / assert
with pytest.raises(RuntimeError, match="boom"):
harness.orchestrator.run_subtask(subtask.id, work=boom)
assert harness.subtasks.get(subtask.id).status is SubTaskStatus.FAILED
assert harness.tasks.get(task.id).status is TaskStatus.FAILED
def test_run_subtasks_creates_runs_and_completes_a_whole_batch(
harness: Harness,
) -> None:
"""run_subtasks fans the parent into one child per item, runs each, and
leaves every child COMPLETE with the parent cascaded to COMPLETE."""
# arrange — a parent task whose coordinator subtask is done
task, coordinator = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
harness.orchestrator.complete_subtask(coordinator.id)
# act
results = harness.orchestrator.run_subtasks(
task.id,
[{"property_id": 1}, {"property_id": 2}, {"property_id": 3}],
work=lambda st: (st.inputs or {})["property_id"] * 10,
)
# assert — results in order, all children COMPLETE, parent COMPLETE
assert results == [10, 20, 30]
children = [s for s in harness.subtasks.list_by_task(task.id) if s.id != coordinator.id]
assert len(children) == 3
assert all(c.status is SubTaskStatus.COMPLETE for c in children)
assert {(c.inputs or {})["property_id"] for c in children} == {1, 2, 3}
assert harness.tasks.get(task.id).status is TaskStatus.COMPLETE
def test_run_subtasks_isolates_a_failing_item_and_continues(
harness: Harness,
) -> None:
"""A raising item is marked FAILED with its error recorded, its siblings still
complete (result None for the failure), and the parent cascades to FAILED."""
# arrange
task, coordinator = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
harness.orchestrator.complete_subtask(coordinator.id)
def _work(st: SubTask) -> int:
pid = (st.inputs or {})["property_id"]
if pid == 2:
raise RuntimeError("property 2 exploded")
return pid
# act — does NOT raise; failure is isolated
results = harness.orchestrator.run_subtasks(
task.id,
[{"property_id": 1}, {"property_id": 2}, {"property_id": 3}],
work=_work,
)
# assert
assert results == [1, None, 3]
children = {
(c.inputs or {})["property_id"]: c
for c in harness.subtasks.list_by_task(task.id)
if c.id != coordinator.id
}
assert children[1].status is SubTaskStatus.COMPLETE
assert children[3].status is SubTaskStatus.COMPLETE
assert children[2].status is SubTaskStatus.FAILED
assert children[2].outputs == {"error": "property 2 exploded"}
# any FAILED child → parent FAILED
assert harness.tasks.get(task.id).status is TaskStatus.FAILED
def test_completing_a_rerun_failed_subtask_unfails_the_task(
harness: Harness,
) -> None:
"""A failed batch is re-run by re-sending its own inputs; when the re-run
completes, the parent Task recomputes to COMPLETE (ADR-0055)."""
# arrange — two children: one failed (task FAILED), the other complete
task, failed_batch = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
healthy_batch = harness.orchestrator.create_child_subtask(task.id)
harness.orchestrator.complete_subtask(healthy_batch.id)
harness.orchestrator.fail_subtask(failed_batch.id, RuntimeError("boom"))
assert harness.tasks.get(task.id).status is TaskStatus.FAILED
# act — the re-run of the failed batch
harness.orchestrator.run_subtask(failed_batch.id, work=lambda: "fixed")
# assert — no failed children remain, so the task un-fails to COMPLETE
task_after = harness.tasks.get(task.id)
assert task_after.status is TaskStatus.COMPLETE
assert task_after.job_completed is not None
def test_task_stays_failed_while_a_failed_subtask_remains(harness: Harness) -> None:
"""Completing a *sibling* does not un-fail the Task — only re-running the
failed SubTask itself can (see the re-run test above)."""
# arrange — two children; fail the first so the task is FAILED
task, coordinator = harness.orchestrator.create_task_with_subtask(
task_source="manual:test"
)
second = harness.orchestrator.create_child_subtask(task.id)
harness.orchestrator.fail_subtask(coordinator.id, RuntimeError("boom"))
assert harness.tasks.get(task.id).status is TaskStatus.FAILED
# act — complete the other child
harness.orchestrator.complete_subtask(second.id)
# assert — task stays FAILED (cascade short-circuited on the terminal state)
assert harness.tasks.get(task.id).status is TaskStatus.FAILED