wehub-resource-sync
76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
"""Cascade kill walks past a single page so >page_size descendants are caught."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from opensquilla.session.manager import SessionManager
|
|
from opensquilla.session.models import SessionNode, SessionStatus
|
|
|
|
|
|
class _PaginatingStorage:
|
|
"""In-memory storage that paginates list_sessions(spawned_by=...) at 100/page."""
|
|
|
|
PAGE_SIZE = 100
|
|
|
|
def __init__(self) -> None:
|
|
self.sessions: dict[str, SessionNode] = {}
|
|
|
|
async def get_session(self, session_key: str):
|
|
return self.sessions.get(session_key)
|
|
|
|
async def upsert_session(self, node: SessionNode) -> None:
|
|
self.sessions[node.session_key] = node
|
|
|
|
async def list_sessions(
|
|
self,
|
|
agent_id=None,
|
|
status=None,
|
|
limit=PAGE_SIZE,
|
|
offset=0,
|
|
spawned_by=None,
|
|
):
|
|
rows = list(self.sessions.values())
|
|
if status is not None:
|
|
rows = [r for r in rows if str(r.status) == str(status)]
|
|
if spawned_by is not None:
|
|
rows = [r for r in rows if r.spawned_by == spawned_by]
|
|
return rows[offset : offset + limit]
|
|
|
|
|
|
def _running_node(session_key: str, *, spawned_by: str | None = None) -> SessionNode:
|
|
return SessionNode(
|
|
session_key=session_key,
|
|
session_id=session_key.replace(":", "-"),
|
|
agent_id="main",
|
|
created_at=1,
|
|
updated_at=1,
|
|
started_at=1,
|
|
status=SessionStatus.RUNNING,
|
|
spawned_by=spawned_by,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cascade_kills_descendants_beyond_first_page() -> None:
|
|
"""A parent with 250 running children — all must end killed after cascade."""
|
|
storage = _PaginatingStorage()
|
|
storage.sessions["agent:main:main"] = _running_node("agent:main:main")
|
|
for i in range(250):
|
|
key = f"agent:main:subagent:{i:04d}"
|
|
storage.sessions[key] = _running_node(key, spawned_by="agent:main:main")
|
|
|
|
class _Registry:
|
|
async def list_agents(self, *, include_builtin: bool = True):
|
|
return [{"id": "main", "enabled": True}]
|
|
|
|
mgr = SessionManager(storage, agent_registry=_Registry()) # type: ignore[arg-type]
|
|
|
|
await mgr.kill_session("agent:main:main")
|
|
|
|
killed = sum(
|
|
1 for s in storage.sessions.values() if s.status == SessionStatus.KILLED
|
|
)
|
|
# Parent + 250 children = 251 total kills.
|
|
assert killed == 251
|