- adapter.py: add completed_at field to C2TaskStatus dataclass - mythic.py: implement get_task() (GraphQL task query) and get_task_output() (response query + decode_response_text concat) - fake.py: deterministic state progression via per-instance call counter; get_task_output raises C2Error until completed - mapping.py: apply_task_to_simulation() idempotent output mapper (mapping_applied anchor prevents double-writes) - migration 0007: add mapping_applied BOOLEAN NOT NULL DEFAULT false to c2_task - c2_task model: mapping_applied column added - api/c2.py: GET /api/simulations/<id>/c2/tasks poll-on-read endpoint; refreshes incomplete tasks from C2, fetches output on completion, applies mapping, skips re-polling for completed tasks; best-effort (C2Error on individual task skipped, returns 200 with stale status) - 51 new tests (396 total); pytest/ruff/mypy all green Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
98 lines
3.1 KiB
Python
98 lines
3.1 KiB
Python
"""Unit tests for apply_task_to_simulation() mapping helper."""
|
|
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
from unittest.mock import MagicMock
|
|
|
|
from backend.app.services.c2.mapping import apply_task_to_simulation
|
|
|
|
|
|
def _make_task(output: str | None = "whoami output", mapping_applied: bool = False) -> MagicMock:
|
|
task = MagicMock()
|
|
task.output = output
|
|
task.mapping_applied = mapping_applied
|
|
return task
|
|
|
|
|
|
def _make_sim(execution_result: str | None = None) -> MagicMock:
|
|
sim = MagicMock()
|
|
sim.execution_result = execution_result
|
|
sim.updated_at = None
|
|
return sim
|
|
|
|
|
|
class TestApplyTaskToSimulation:
|
|
def test_appends_output_to_empty_simulation(self):
|
|
task = _make_task(output="whoami output")
|
|
sim = _make_sim(execution_result=None)
|
|
|
|
apply_task_to_simulation(task, sim)
|
|
|
|
assert sim.execution_result == "whoami output"
|
|
assert task.mapping_applied is True
|
|
|
|
def test_appends_with_newline_separator(self):
|
|
task = _make_task(output="second result")
|
|
sim = _make_sim(execution_result="first result")
|
|
|
|
apply_task_to_simulation(task, sim)
|
|
|
|
assert sim.execution_result == "first result\nsecond result"
|
|
|
|
def test_idempotent_when_already_applied(self):
|
|
task = _make_task(output="some output", mapping_applied=True)
|
|
sim = _make_sim(execution_result="existing")
|
|
|
|
apply_task_to_simulation(task, sim)
|
|
|
|
# execution_result must not be modified.
|
|
assert sim.execution_result == "existing"
|
|
|
|
def test_no_op_when_output_is_empty_string(self):
|
|
task = _make_task(output="")
|
|
sim = _make_sim(execution_result="existing")
|
|
|
|
apply_task_to_simulation(task, sim)
|
|
|
|
assert sim.execution_result == "existing"
|
|
# Still marks mapping_applied so we don't revisit it.
|
|
assert task.mapping_applied is True
|
|
|
|
def test_no_op_when_output_is_none(self):
|
|
task = _make_task(output=None)
|
|
sim = _make_sim(execution_result="existing")
|
|
|
|
apply_task_to_simulation(task, sim)
|
|
|
|
assert sim.execution_result == "existing"
|
|
assert task.mapping_applied is True
|
|
|
|
def test_strips_trailing_newlines_from_existing(self):
|
|
"""Existing execution_result with trailing newlines should not cause double blank lines."""
|
|
task = _make_task(output="new output")
|
|
sim = _make_sim(execution_result="old output\n\n")
|
|
|
|
apply_task_to_simulation(task, sim)
|
|
|
|
assert sim.execution_result == "old output\nnew output"
|
|
|
|
def test_updated_at_is_set_on_sim(self):
|
|
task = _make_task(output="something")
|
|
sim = _make_sim(execution_result=None)
|
|
before = datetime.now(UTC)
|
|
|
|
apply_task_to_simulation(task, sim)
|
|
|
|
assert sim.updated_at is not None
|
|
assert sim.updated_at >= before
|
|
|
|
def test_multiple_tasks_accumulate(self):
|
|
sim = _make_sim(execution_result=None)
|
|
tasks = [_make_task(output=f"result {i}") for i in range(3)]
|
|
|
|
for t in tasks:
|
|
apply_task_to_simulation(t, sim)
|
|
|
|
lines = sim.execution_result.split("\n")
|
|
assert lines == ["result 0", "result 1", "result 2"]
|