295 lines
9.2 KiB
Python
295 lines
9.2 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
from datetime import UTC, datetime
|
|
from decimal import Decimal
|
|
from types import SimpleNamespace
|
|
from uuid import RFC_4122, UUID
|
|
|
|
import pytest
|
|
|
|
from src.db.awooop_models import AwoooPRunIdempotency, AwoooPRunState
|
|
from src.services import platform_runtime
|
|
from src.services.platform_operator_service import _run_list_item
|
|
from src.services.runtime_correlation import (
|
|
canonical_traceparent,
|
|
canonical_work_item_id,
|
|
correlation_for_run,
|
|
correlation_readback_for_run,
|
|
is_canonical_traceparent,
|
|
new_uuid7,
|
|
)
|
|
|
|
|
|
def test_uuid7_has_rfc_version_variant_and_timestamp_layout() -> None:
|
|
timestamp_ms = 1_753_155_200_123
|
|
generated = new_uuid7(timestamp_ms=timestamp_ms, entropy=(1 << 74) - 1)
|
|
|
|
assert generated.version == 7
|
|
assert generated.variant == RFC_4122
|
|
assert generated.int >> 80 == timestamp_ms
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("timestamp_ms", "entropy"),
|
|
[(-1, 0), (1 << 48, 0), (0, -1), (0, 1 << 74)],
|
|
)
|
|
def test_uuid7_rejects_out_of_range_inputs(
|
|
timestamp_ms: int,
|
|
entropy: int,
|
|
) -> None:
|
|
with pytest.raises(ValueError):
|
|
new_uuid7(timestamp_ms=timestamp_ms, entropy=entropy)
|
|
|
|
|
|
def test_runtime_correlation_is_stable_and_bound_to_run() -> None:
|
|
run_id = UUID("019f88d5-f6c9-7a18-b764-df0df0cbe222")
|
|
correlation = correlation_for_run("awoooi", run_id)
|
|
|
|
assert correlation.run_id == run_id
|
|
assert correlation.trace_id == canonical_traceparent(run_id)
|
|
assert correlation.work_item_id == f"platform-run:awoooi:{run_id}"
|
|
assert is_canonical_traceparent(correlation.trace_id, run_id=run_id) is True
|
|
assert (
|
|
is_canonical_traceparent(
|
|
correlation.trace_id,
|
|
run_id=UUID("019f88d5-f6c9-7a18-b764-df0df0cbe223"),
|
|
)
|
|
is False
|
|
)
|
|
|
|
|
|
def test_runtime_correlation_rejects_invalid_project_id() -> None:
|
|
with pytest.raises(ValueError, match="project_id_invalid"):
|
|
canonical_work_item_id(
|
|
"awoooi:forged",
|
|
UUID("019f88d5-f6c9-7a18-b764-df0df0cbe222"),
|
|
)
|
|
|
|
|
|
def test_runtime_correlation_readback_preserves_legacy_trace_truth() -> None:
|
|
run_id = UUID("019f88d5-f6c9-7a18-b764-df0df0cbe222")
|
|
|
|
canonical = correlation_readback_for_run(
|
|
"awoooi",
|
|
run_id,
|
|
canonical_traceparent(run_id),
|
|
)
|
|
legacy = correlation_readback_for_run("awoooi", run_id, "legacy-trace")
|
|
|
|
assert canonical == {
|
|
"trace_id": canonical_traceparent(run_id),
|
|
"work_item_id": f"platform-run:awoooi:{run_id}",
|
|
"correlation_status": "canonical",
|
|
}
|
|
assert legacy == {
|
|
"trace_id": "legacy-trace",
|
|
"work_item_id": f"platform-run:awoooi:{run_id}",
|
|
"correlation_status": "legacy_or_invalid_trace",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_status_exposes_canonical_identity_readback(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
run_id = UUID("019f88d5-f6c9-7a18-b764-df0df0cbe222")
|
|
trace_id = canonical_traceparent(run_id)
|
|
row = SimpleNamespace(
|
|
run_id=run_id,
|
|
project_id="awoooi",
|
|
agent_id="runtime-agent",
|
|
state="running",
|
|
is_shadow=True,
|
|
trace_id=trace_id,
|
|
attempt_count=1,
|
|
cost_usd=Decimal("0.0000"),
|
|
step_count=2,
|
|
error_code=None,
|
|
created_at=datetime(2026, 7, 22, tzinfo=UTC),
|
|
started_at=datetime(2026, 7, 22, tzinfo=UTC),
|
|
completed_at=None,
|
|
timeout_at=None,
|
|
)
|
|
|
|
class _Result:
|
|
def scalar_one_or_none(self) -> SimpleNamespace:
|
|
return row
|
|
|
|
class _Db:
|
|
async def execute(self, _statement: object) -> _Result:
|
|
return _Result()
|
|
|
|
@asynccontextmanager
|
|
async def fake_db_context(project_id: str):
|
|
assert project_id == "awoooi"
|
|
yield _Db()
|
|
|
|
monkeypatch.setattr(platform_runtime, "get_db_context", fake_db_context)
|
|
|
|
payload = await platform_runtime.get_run_status(run_id, "awoooi")
|
|
|
|
assert payload is not None
|
|
assert payload["run_id"] == str(run_id)
|
|
assert payload["trace_id"] == trace_id
|
|
assert payload["work_item_id"] == f"platform-run:awoooi:{run_id}"
|
|
assert payload["correlation_status"] == "canonical"
|
|
list_item = _run_list_item(
|
|
row,
|
|
remediation_summary=None,
|
|
callback_reply_summary=None,
|
|
)
|
|
assert {
|
|
key: list_item[key]
|
|
for key in ("trace_id", "work_item_id", "correlation_status")
|
|
} == {
|
|
key: payload[key]
|
|
for key in ("trace_id", "work_item_id", "correlation_status")
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_status_marks_legacy_unbound_trace_without_rewriting_it(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
run_id = UUID("019f88d5-f6c9-7a18-b764-df0df0cbe222")
|
|
row = SimpleNamespace(
|
|
run_id=run_id,
|
|
project_id="awoooi",
|
|
agent_id="legacy-agent",
|
|
state="completed",
|
|
is_shadow=True,
|
|
trace_id="legacy-trace",
|
|
attempt_count=1,
|
|
cost_usd=Decimal("0.0000"),
|
|
step_count=0,
|
|
error_code=None,
|
|
created_at=None,
|
|
started_at=None,
|
|
completed_at=None,
|
|
)
|
|
|
|
class _Result:
|
|
def scalar_one_or_none(self) -> SimpleNamespace:
|
|
return row
|
|
|
|
class _Db:
|
|
async def execute(self, _statement: object) -> _Result:
|
|
return _Result()
|
|
|
|
@asynccontextmanager
|
|
async def fake_db_context(_project_id: str):
|
|
yield _Db()
|
|
|
|
monkeypatch.setattr(platform_runtime, "get_db_context", fake_db_context)
|
|
|
|
payload = await platform_runtime.get_run_status(run_id, "awoooi")
|
|
|
|
assert payload is not None
|
|
assert payload["trace_id"] == "legacy-trace"
|
|
assert payload["correlation_status"] == "legacy_or_invalid_trace"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_provider_event_creation_returns_one_durable_winner(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
class _Result:
|
|
def __init__(self, row: tuple[UUID] | None = None) -> None:
|
|
self._row = row
|
|
|
|
def fetchone(self) -> tuple[UUID] | None:
|
|
return self._row
|
|
|
|
class _Store:
|
|
def __init__(self) -> None:
|
|
self.claim_lock = asyncio.Lock()
|
|
self.mapping: UUID | None = None
|
|
self.runs: dict[UUID, AwoooPRunState] = {}
|
|
self.lock_claim_count = 0
|
|
|
|
store = _Store()
|
|
|
|
class _Db:
|
|
def __init__(self) -> None:
|
|
self.lock_held = False
|
|
self.pending_run: AwoooPRunState | None = None
|
|
self.pending_mapping: AwoooPRunIdempotency | None = None
|
|
|
|
async def execute(
|
|
self,
|
|
statement: object,
|
|
_params: dict[str, str] | None = None,
|
|
) -> _Result:
|
|
sql = str(statement)
|
|
if "pg_advisory_xact_lock" in sql:
|
|
await store.claim_lock.acquire()
|
|
self.lock_held = True
|
|
store.lock_claim_count += 1
|
|
await asyncio.sleep(0)
|
|
return _Result()
|
|
if "awooop_run_idempotency" in sql:
|
|
return _Result(
|
|
(store.mapping,) if store.mapping is not None else None
|
|
)
|
|
raise AssertionError(f"unexpected statement: {sql}")
|
|
|
|
def add(self, row: object) -> None:
|
|
if isinstance(row, AwoooPRunState):
|
|
self.pending_run = row
|
|
elif isinstance(row, AwoooPRunIdempotency):
|
|
self.pending_mapping = row
|
|
else:
|
|
raise AssertionError(f"unexpected row: {type(row).__name__}")
|
|
|
|
async def flush(self) -> None:
|
|
assert self.lock_held is True
|
|
assert self.pending_run is not None
|
|
if self.pending_mapping is None:
|
|
assert store.runs == {}
|
|
store.runs[self.pending_run.run_id] = self.pending_run
|
|
return
|
|
assert store.mapping is None
|
|
assert self.pending_mapping.run_id in store.runs
|
|
store.mapping = self.pending_mapping.run_id
|
|
|
|
@asynccontextmanager
|
|
async def fake_db_context(project_id: str):
|
|
assert project_id == "awoooi"
|
|
db = _Db()
|
|
try:
|
|
yield db
|
|
finally:
|
|
if db.lock_held:
|
|
store.claim_lock.release()
|
|
|
|
async def fake_cache(*_args: object) -> None:
|
|
return None
|
|
|
|
monkeypatch.setattr(platform_runtime, "get_db_context", fake_db_context)
|
|
monkeypatch.setattr(platform_runtime, "_cache_idempotency", fake_cache)
|
|
|
|
first, second = await asyncio.gather(
|
|
platform_runtime.create_run(
|
|
project_id="awoooi",
|
|
agent_id="runtime-agent",
|
|
trigger_type="channel_event",
|
|
channel_type="telegram",
|
|
provider_event_id="same-provider-event",
|
|
),
|
|
platform_runtime.create_run(
|
|
project_id="awoooi",
|
|
agent_id="runtime-agent",
|
|
trigger_type="channel_event",
|
|
channel_type="telegram",
|
|
provider_event_id="same-provider-event",
|
|
),
|
|
)
|
|
|
|
assert first[0] == second[0] == store.mapping
|
|
assert sorted((first[1], second[1])) == [False, True]
|
|
assert list(store.runs) == [store.mapping]
|
|
assert next(iter(store.runs.values())).timeout_at.tzinfo is None
|
|
assert store.lock_claim_count == 2
|