wehub-resource-sync
134 lines
4.4 KiB
Python
134 lines
4.4 KiB
Python
"""Tests for C3 fix: no split-brain lock on rapid re-enqueue after terminal.
|
|
|
|
AC-C3-1: _session_locks is never popped in _mark_terminal.
|
|
AC-C3-2: rapid enqueue -> terminal -> re-enqueue for same session_key never
|
|
allows two tasks to run concurrently (max_concurrent_per_session == 1).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from opensquilla.gateway.routing import RouteEnvelope, SourceKind
|
|
from opensquilla.gateway.task_runtime import TaskRuntime
|
|
from opensquilla.session.models import AgentTaskRecord
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _make_envelope(session_key: str = "agent-1::sess-1") -> RouteEnvelope:
|
|
return RouteEnvelope(
|
|
source_kind=SourceKind.WEB,
|
|
source_name="test",
|
|
agent_id="agent-1",
|
|
session_key=session_key,
|
|
input_provenance={"kind": "test"},
|
|
)
|
|
|
|
|
|
def _make_storage() -> Any:
|
|
storage = MagicMock()
|
|
task_db: dict[str, AgentTaskRecord] = {}
|
|
|
|
async def create(record: AgentTaskRecord) -> None:
|
|
task_db[record.task_id] = record
|
|
|
|
async def update(task_id: str, **kwargs: Any) -> None:
|
|
rec = task_db.get(task_id)
|
|
if rec is None:
|
|
return
|
|
for k, v in kwargs.items():
|
|
if hasattr(rec, k):
|
|
object.__setattr__(rec, k, v)
|
|
|
|
async def get(task_id: str) -> AgentTaskRecord | None:
|
|
return task_db.get(task_id)
|
|
|
|
async def list_tasks(**_: Any) -> list[AgentTaskRecord]:
|
|
return list(task_db.values())
|
|
|
|
storage.create_agent_task = create
|
|
storage.update_agent_task = update
|
|
storage.get_agent_task = get
|
|
storage.list_agent_tasks = list_tasks
|
|
return storage
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# AC-C3-1: _session_locks never popped in _mark_terminal
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_session_locks_never_popped_at_terminal() -> None:
|
|
"""After a task reaches terminal state, _session_locks still contains the key."""
|
|
async def _instant(_run: Any) -> None:
|
|
pass
|
|
|
|
rt = TaskRuntime(
|
|
storage=_make_storage(),
|
|
turn_handler=_instant,
|
|
max_concurrency=4,
|
|
)
|
|
env = _make_envelope("agent-1::sess-c3")
|
|
handle = await rt.enqueue(env, "msg")
|
|
await rt.wait(handle.task_id, timeout=5.0)
|
|
|
|
# Lock must still be present — C3 fix ensures we never pop it.
|
|
assert env.session_key in rt._session_locks, (
|
|
"_session_locks should retain the entry after terminal (C3 fix)"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# AC-C3-2: no split-brain under rapid enqueue -> terminal -> re-enqueue
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rapid_enqueue_after_terminal_no_split_brain() -> None:
|
|
"""Loop 100 times: enqueue -> wait for terminal -> immediately re-enqueue.
|
|
|
|
At no point should two tasks for the same session run concurrently.
|
|
max_concurrent_per_session must always be 1.
|
|
"""
|
|
iterations = 100
|
|
session_key = "agent-1::sess-rapid"
|
|
env = _make_envelope(session_key)
|
|
|
|
concurrent_count = 0
|
|
max_concurrent = 0
|
|
count_lock = asyncio.Lock()
|
|
|
|
async def _handler(_run: Any) -> None:
|
|
nonlocal concurrent_count, max_concurrent
|
|
async with count_lock:
|
|
concurrent_count += 1
|
|
if concurrent_count > max_concurrent:
|
|
max_concurrent = concurrent_count
|
|
# Small yield to allow other tasks to slip in if the lock is broken.
|
|
await asyncio.sleep(0)
|
|
async with count_lock:
|
|
concurrent_count -= 1
|
|
|
|
rt = TaskRuntime(
|
|
storage=_make_storage(),
|
|
turn_handler=_handler,
|
|
max_concurrency=4,
|
|
max_pending_per_session=None,
|
|
)
|
|
|
|
for _ in range(iterations):
|
|
handle = await rt.enqueue(env, "msg")
|
|
await rt.wait(handle.task_id, timeout=5.0)
|
|
# Yield to event loop so any in-flight concurrent task could manifest.
|
|
await asyncio.sleep(0)
|
|
|
|
assert max_concurrent == 1, (
|
|
f"Split-brain detected: max concurrent tasks for same session = {max_concurrent} "
|
|
f"(expected 1)"
|
|
)
|