feat(backend): c2 poll-on-read + output mapping (sprint 8 M3)
- 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>
This commit is contained in:
124
backend/tests/test_migration_0007_c2.py
Normal file
124
backend/tests/test_migration_0007_c2.py
Normal file
@@ -0,0 +1,124 @@
|
||||
"""Migration round-trip test for 0007_c2_task_mapping_applied.
|
||||
|
||||
Verifies that upgrade() adds the mapping_applied column and downgrade() removes it.
|
||||
Uses the resolved-path pattern per lessons.md Sprint 4.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
|
||||
from alembic.operations import Operations
|
||||
from alembic.runtime.migration import MigrationContext
|
||||
from sqlalchemy import create_engine, inspect, text
|
||||
|
||||
|
||||
def _load_migration(name: str):
|
||||
versions_dir = Path(__file__).resolve().parent.parent / "migrations" / "versions"
|
||||
path = versions_dir / name
|
||||
spec = importlib.util.spec_from_file_location(name.removesuffix(".py"), path)
|
||||
assert spec and spec.loader
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod) # type: ignore[union-attr]
|
||||
return mod
|
||||
|
||||
|
||||
def _fresh_engine_with_c2_task():
|
||||
"""In-memory SQLite with c2_task already created (as left by 0006 upgrade)."""
|
||||
engine = create_engine("sqlite:///:memory:")
|
||||
with engine.begin() as conn:
|
||||
conn.execute(text("""
|
||||
CREATE TABLE c2_task (
|
||||
id INTEGER PRIMARY KEY,
|
||||
simulation_id INTEGER NOT NULL,
|
||||
mythic_task_display_id INTEGER NOT NULL,
|
||||
callback_display_id INTEGER NOT NULL,
|
||||
command TEXT NOT NULL,
|
||||
params TEXT,
|
||||
status TEXT NOT NULL,
|
||||
completed BOOLEAN NOT NULL DEFAULT 0,
|
||||
output TEXT,
|
||||
source TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL,
|
||||
completed_at DATETIME
|
||||
)
|
||||
"""))
|
||||
return engine
|
||||
|
||||
|
||||
def _run_upgrade(engine, migration_mod):
|
||||
with engine.begin() as conn:
|
||||
ctx = MigrationContext.configure(conn)
|
||||
ops = Operations(ctx)
|
||||
ops._install_proxy() # type: ignore[attr-defined]
|
||||
try:
|
||||
migration_mod.upgrade()
|
||||
finally:
|
||||
ops._remove_proxy() # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def _run_downgrade(engine, migration_mod):
|
||||
with engine.begin() as conn:
|
||||
ctx = MigrationContext.configure(conn)
|
||||
ops = Operations(ctx)
|
||||
ops._install_proxy() # type: ignore[attr-defined]
|
||||
try:
|
||||
migration_mod.downgrade()
|
||||
finally:
|
||||
ops._remove_proxy() # type: ignore[attr-defined]
|
||||
|
||||
|
||||
class TestMigration0007Upgrade:
|
||||
def test_mapping_applied_column_added(self):
|
||||
engine = _fresh_engine_with_c2_task()
|
||||
mod = _load_migration("0007_c2_task_mapping_applied.py")
|
||||
_run_upgrade(engine, mod)
|
||||
|
||||
insp = inspect(engine)
|
||||
cols = {c["name"] for c in insp.get_columns("c2_task")}
|
||||
assert "mapping_applied" in cols
|
||||
|
||||
def test_mapping_applied_defaults_to_false(self):
|
||||
engine = _fresh_engine_with_c2_task()
|
||||
mod = _load_migration("0007_c2_task_mapping_applied.py")
|
||||
|
||||
# Insert a row before upgrading (no mapping_applied column yet).
|
||||
with engine.begin() as conn:
|
||||
conn.execute(text(
|
||||
"INSERT INTO c2_task "
|
||||
"(simulation_id, mythic_task_display_id, callback_display_id, "
|
||||
"command, status, completed, source, created_at) "
|
||||
"VALUES (1, 1000, 1, 'whoami', 'submitted', 0, 'mimic', '2026-01-01')"
|
||||
))
|
||||
|
||||
_run_upgrade(engine, mod)
|
||||
|
||||
with engine.begin() as conn:
|
||||
row = conn.execute(
|
||||
text("SELECT mapping_applied FROM c2_task WHERE id = 1")
|
||||
).fetchone()
|
||||
assert row is not None
|
||||
# SQLite stores booleans as 0/1.
|
||||
assert row[0] == 0 or row[0] is False
|
||||
|
||||
|
||||
class TestMigration0007Downgrade:
|
||||
def test_downgrade_removes_mapping_applied(self):
|
||||
engine = _fresh_engine_with_c2_task()
|
||||
mod = _load_migration("0007_c2_task_mapping_applied.py")
|
||||
_run_upgrade(engine, mod)
|
||||
_run_downgrade(engine, mod)
|
||||
|
||||
insp = inspect(engine)
|
||||
cols = {c["name"] for c in insp.get_columns("c2_task")}
|
||||
assert "mapping_applied" not in cols
|
||||
|
||||
def test_downgrade_does_not_drop_other_columns(self):
|
||||
engine = _fresh_engine_with_c2_task()
|
||||
mod = _load_migration("0007_c2_task_mapping_applied.py")
|
||||
_run_upgrade(engine, mod)
|
||||
_run_downgrade(engine, mod)
|
||||
|
||||
insp = inspect(engine)
|
||||
cols = {c["name"] for c in insp.get_columns("c2_task")}
|
||||
assert {"id", "simulation_id", "command", "status", "completed"} <= cols
|
||||
Reference in New Issue
Block a user