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