273 lines
10 KiB
Python
273 lines
10 KiB
Python
"""In-memory task/run fixtures exercising the production repository contract."""
|
|
from contextlib import asynccontextmanager
|
|
from copy import deepcopy
|
|
from datetime import datetime, timedelta
|
|
from uuid import uuid4
|
|
from app.auth.project_dependencies import ProjectContext
|
|
from app.domain.schemas.operational_task import OperationalTaskSummary, OperationalTaskRun
|
|
|
|
def _context(user_id=None, *, system_role="user"):
|
|
return ProjectContext(
|
|
project_id=uuid4(), project_code="test", user_id=user_id or uuid4(),
|
|
project_role="member", system_role=system_role,
|
|
)
|
|
|
|
def _task_document(context, now, *, lease=None):
|
|
task_id = uuid4()
|
|
return OperationalTaskSummary(
|
|
execution_scope="production",
|
|
project_id=context.project_id,
|
|
task_id=task_id,
|
|
task_type="event_investigation",
|
|
workflow="water_quality_complaint",
|
|
status="running",
|
|
priority="normal",
|
|
revision=2,
|
|
execution_generation=1,
|
|
execution_request_id=uuid4(),
|
|
title="调查",
|
|
summary="验证状态机",
|
|
source={"kind": "operator", "event_kind": "water_quality_complaint"},
|
|
trigger={"kind": "manual", "occurred_at": now, "detail": "test"},
|
|
scheduled_for=None,
|
|
agent_session_id="session-1",
|
|
agent_lease_expires_at=lease,
|
|
current_plan_revision=None,
|
|
current_report_revision=None,
|
|
created_at=now,
|
|
updated_at=now,
|
|
).model_dump(mode="json")
|
|
|
|
def _active_run(task, now):
|
|
return OperationalTaskRun(
|
|
run_id=uuid4(), task_id=task["task_id"], generation=task["execution_generation"],
|
|
execution_request_id=task["execution_request_id"], status="running", revision=2,
|
|
agent_session_id="session-1", lease_expires_at=now + timedelta(minutes=1),
|
|
heartbeat_at=now, started_at=now, dispatch_attempts=0,
|
|
next_dispatch_at=now, created_at=now, updated_at=now,
|
|
).model_dump(mode="json")
|
|
|
|
class TaskRepository:
|
|
def __init__(self, task, run=None, *, action=None):
|
|
self.task = task
|
|
self.action = action
|
|
self.requests = {}
|
|
self.evidence = {}
|
|
self.feedback = {}
|
|
self.plans = []
|
|
self.reports = []
|
|
self.progress = {}
|
|
self.scada_presentations = {}
|
|
self.notifications = []
|
|
self.locks = []
|
|
self.access_requests = {}
|
|
self.events = []
|
|
self.questions = {}
|
|
now = datetime.fromisoformat(task["created_at"])
|
|
seeded = _active_run(task, now) if run is None else run
|
|
if run is None:
|
|
seeded.update(
|
|
status=task["status"] if task["status"] in {"queued", "running"} else "succeeded",
|
|
agent_session_id=task.get("agent_session_id"),
|
|
lease_expires_at=task.get("agent_lease_expires_at"),
|
|
)
|
|
self.run = seeded
|
|
self.runs = {seeded["run_id"]: seeded}
|
|
|
|
@asynccontextmanager
|
|
async def transaction(self):
|
|
yield self
|
|
|
|
async def lock(self, key):
|
|
self.locks.append(key)
|
|
|
|
async def get_request(self, scope, request_id):
|
|
return self.requests.get((scope, request_id))
|
|
|
|
async def save_request(self, scope, request_id, request_digest, response):
|
|
self.requests[(scope, request_id)] = {
|
|
"request_digest": request_digest,
|
|
"response": response,
|
|
}
|
|
|
|
async def get_task(self, task_id):
|
|
return self.task if task_id == self.task["task_id"] else None
|
|
|
|
async def update_task(self, task):
|
|
self.task = task
|
|
|
|
async def get_latest_plan(self, _task_id):
|
|
return self.plans[-1] if self.plans else None
|
|
|
|
async def get_latest_report(self, _task_id):
|
|
return self.reports[-1] if self.reports else None
|
|
|
|
async def get_latest_progress(self, task_id, generation):
|
|
return self.progress.get((task_id, generation))
|
|
|
|
async def get_scada_presentation(self, task_id, run_id):
|
|
return deepcopy(self.scada_presentations.get((task_id, run_id)))
|
|
|
|
async def save_scada_presentation(self, task_id, run_id, document, now):
|
|
self.scada_presentations[(task_id, run_id)] = deepcopy(document)
|
|
|
|
async def get_latest_graph_progress(self, task_id):
|
|
snapshots = [value for (stored_id, _), value in self.progress.items()
|
|
if stored_id == task_id
|
|
and (value["progress"].get("diagram") or {}).get("edges") is not None]
|
|
return max(snapshots, key=lambda row: (row["execution_generation"], row["revision"])) if snapshots else None
|
|
|
|
async def get_latest_diagram(self, task_id):
|
|
diagrams = [(row["published_at"], row["progress"].get("diagram"))
|
|
for (stored_id, _), row in self.progress.items() if stored_id == task_id]
|
|
diagrams.extend((row["published_at"], row["report"].get("diagram")) for row in self.reports)
|
|
return next((diagram for _, diagram in sorted(diagrams, key=lambda item: item[0], reverse=True)
|
|
if diagram is not None), None)
|
|
|
|
async def save_progress(self, task_id, snapshot):
|
|
self.progress[(task_id, snapshot["execution_generation"])] = snapshot
|
|
|
|
async def get_report(self, _task_id, revision):
|
|
return next((item for item in self.reports if item["revision"] == revision), None)
|
|
|
|
async def list_reports(self, _task_id, limit, offset):
|
|
return list(reversed(self.reports))[offset:offset + limit]
|
|
|
|
async def insert_plan(self, _task_id, snapshot):
|
|
self.plans.append(snapshot)
|
|
|
|
async def insert_report(self, _task_id, snapshot):
|
|
self.reports.append(snapshot)
|
|
|
|
async def list_evidence(self, _task_id):
|
|
return list(self.evidence.values())
|
|
|
|
async def get_evidence(self, _task_id, evidence_id):
|
|
return self.evidence.get(evidence_id)
|
|
|
|
async def insert_evidence(self, _task_id, evidence):
|
|
self.evidence[evidence["evidence_id"]] = evidence
|
|
|
|
async def list_access_requests(self, _task_id):
|
|
return list(self.access_requests.values())
|
|
|
|
async def get_access_request(self, _task_id, access_request_id):
|
|
return self.access_requests.get(access_request_id)
|
|
|
|
async def insert_access_request(self, _task_id, document):
|
|
self.access_requests[document["access_request_id"]] = document
|
|
|
|
async def update_access_request(self, _task_id, document, expected_version):
|
|
assert document["version"] == expected_version + 1
|
|
self.access_requests[document["access_request_id"]] = document
|
|
|
|
async def list_actions(self, _task_id):
|
|
return [self.action] if self.action else []
|
|
|
|
async def get_action(self, _task_id, action_id):
|
|
if self.action and self.action["action_id"] == action_id:
|
|
return self.action
|
|
return None
|
|
|
|
async def update_action(self, _task_id, action):
|
|
self.action = action
|
|
|
|
async def insert_action(self, _task_id, action):
|
|
self.action = action
|
|
|
|
async def insert_notification(self, notification):
|
|
self.notifications.append(notification)
|
|
|
|
async def resolve_notifications(
|
|
self, task_id, *, kinds, subject_type, subject_id, resolved_at, resolution_reason
|
|
):
|
|
resolved = []
|
|
for notification in self.notifications:
|
|
if (
|
|
notification["task_id"] == task_id
|
|
and notification["kind"] in kinds
|
|
and notification.get("lifecycle_status", "active") == "active"
|
|
and notification.get("subject_type") == subject_type
|
|
and notification.get("subject_id") == subject_id
|
|
):
|
|
notification.update(
|
|
lifecycle_status="resolved",
|
|
resolved_at=resolved_at.isoformat(),
|
|
resolution_reason=resolution_reason,
|
|
)
|
|
resolved.append(notification.copy())
|
|
return resolved
|
|
|
|
async def get_feedback(self, task_id, source_event_id):
|
|
return self.feedback.get((task_id, source_event_id))
|
|
|
|
async def upsert_feedback(self, task_id, feedback, accepted_at):
|
|
self.feedback[(task_id, feedback["source_event_id"])] = {
|
|
"source_revision": feedback["source_revision"],
|
|
"document": feedback,
|
|
"accepted_at": accepted_at,
|
|
}
|
|
|
|
async def get_active_run(self, task_id):
|
|
return next((run for run in sorted(self.runs.values(), key=lambda item: item["generation"], reverse=True)
|
|
if run["task_id"] == task_id and run["status"] in {"queued", "claimed", "running"}), None)
|
|
|
|
async def get_latest_run(self, task_id):
|
|
return next((run for run in sorted(self.runs.values(), key=lambda item: item["generation"], reverse=True) if run["task_id"] == task_id), None)
|
|
|
|
async def get_run(self, run_id):
|
|
return self.runs.get(run_id)
|
|
|
|
async def require_execution_storage(self):
|
|
return None
|
|
|
|
async def update_run(self, run, expected_revision):
|
|
assert run["revision"] == expected_revision + 1
|
|
self.runs[run["run_id"]] = run
|
|
self.run = max(self.runs.values(), key=lambda item: item["generation"])
|
|
|
|
async def insert_run(self, run):
|
|
self.runs[run["run_id"]] = run
|
|
self.run = max(self.runs.values(), key=lambda item: item["generation"])
|
|
|
|
async def append_event(self, **event):
|
|
self.events.append(event)
|
|
return len(self.events)
|
|
|
|
async def current_event_cursor(self):
|
|
return len(self.events)
|
|
|
|
async def list_presentation_events(self, task_id, after, through, limit):
|
|
return [{"event_seq": index, **event} for index, event in enumerate(self.events, 1)
|
|
if after < index <= through and event["topic"] == "operational.presentation.requested"
|
|
and event["aggregate_id"] == task_id][:limit]
|
|
async def access_requests_available(self):
|
|
return True
|
|
|
|
async def fail_run_messages(self, run_id, reason, now):
|
|
return None
|
|
|
|
async def try_lock(self, key):
|
|
await self.lock(key)
|
|
return True
|
|
|
|
async def pending_calculation_jobs(self, task_id, run_id, execution_request_id):
|
|
return []
|
|
|
|
async def pending_run_messages(self, run_id):
|
|
return []
|
|
|
|
async def list_investigation_alerts(self, task_id):
|
|
return []
|
|
|
|
async def list_questions(self, task_id):
|
|
return [deepcopy(q) for q in self.questions.values() if q["task_id"] == task_id]
|
|
|
|
async def get_question(self, task_id, question_id):
|
|
q = self.questions.get(question_id)
|
|
return deepcopy(q) if q and q["task_id"] == task_id else None
|
|
|
|
async def save_question(self, task_id, document):
|
|
assert document["task_id"] == task_id
|
|
self.questions[document["question_request_id"]] = deepcopy(document)
|