87 lines
4.0 KiB
Python
87 lines
4.0 KiB
Python
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
from app.gateway.worker_pool.admission import Admission, bounded_worker_run, run_with_worker_admission
|
|
|
|
|
|
class Bridge:
|
|
def __init__(self): self.events = []
|
|
async def publish(self, run_id, event, data): self.events.append((run_id, event, data))
|
|
async def publish_end(self, run_id): self.events.append((run_id, "end", None))
|
|
async def cleanup(self, *args, **kwargs): pass
|
|
|
|
|
|
class Manager:
|
|
def __init__(self): self.records = {}
|
|
async def set_status(self, run_id, status, **kwargs): self.records[run_id].status = status
|
|
|
|
|
|
async def prepare(monkeypatch, capacity=4, queue=128, timeout=10):
|
|
async def owner(thread): return "owner"
|
|
monkeypatch.setattr("deerflow.runtime.thread_paths.aresolve_path_user_id", owner)
|
|
gate = Admission(capacity, queue, timeout); bridge = Bridge(); manager = Manager()
|
|
app = SimpleNamespace(state=SimpleNamespace(worker_admission=gate, worker_deferred_admission=True, stream_bridge=bridge, run_manager=manager))
|
|
return gate, bridge, manager, SimpleNamespace(app=app)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_70_streams_get_real_run_metadata_before_waiting_for_capacity(monkeypatch):
|
|
gate, bridge, manager, request = await prepare(monkeypatch)
|
|
release = asyncio.Event(); peak = 0
|
|
async def runner(**kwargs):
|
|
nonlocal peak
|
|
peak = max(peak, gate.active)
|
|
await release.wait()
|
|
kwargs["record"].status = "success"
|
|
@bounded_worker_run
|
|
async def start(body, thread_id, request):
|
|
record = SimpleNamespace(run_id=thread_id, thread_id=thread_id, status="pending")
|
|
manager.records[thread_id] = record
|
|
record.task = asyncio.create_task(run_with_worker_admission(runner, bridge=bridge, run_manager=manager, record=record))
|
|
return record
|
|
records = await asyncio.wait_for(asyncio.gather(*(start(None, str(i), request) for i in range(70))), 1)
|
|
await asyncio.sleep(.02)
|
|
assert len([e for e in bridge.events if e[1] == "metadata"]) == 70
|
|
assert gate.active == 4 and gate.waiting == 66
|
|
assert all(not r.task.done() for r in records)
|
|
release.set(); await asyncio.gather(*(r.task for r in records)); await asyncio.sleep(0)
|
|
assert peak == 4 and gate.active == gate.waiting == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_queued_cancel_timeout_and_overflow_finish_stream_without_capacity_leak(monkeypatch):
|
|
gate, bridge, manager, request = await prepare(monkeypatch, capacity=1, queue=2, timeout=.04)
|
|
release = asyncio.Event()
|
|
async def runner(**kwargs): await release.wait()
|
|
@bounded_worker_run
|
|
async def start(body, thread_id, request):
|
|
record = SimpleNamespace(run_id=thread_id, thread_id=thread_id, status="pending")
|
|
manager.records[thread_id] = record
|
|
record.task = asyncio.create_task(run_with_worker_admission(runner, bridge=bridge, run_manager=manager, record=record))
|
|
return record
|
|
first = await start(None, "one", request)
|
|
second = await start(None, "two", request)
|
|
third = await start(None, "three", request)
|
|
with pytest.raises(HTTPException) as exc: await start(None, "overflow", request)
|
|
assert exc.value.status_code == 429
|
|
await asyncio.sleep(.01); second.task.cancel()
|
|
await asyncio.gather(second.task, third.task, return_exceptions=True)
|
|
assert manager.records["two"].status.value == "interrupted"
|
|
assert manager.records["three"].status.value == "error"
|
|
assert any(e[0] == "three" and e[1] == "end" for e in bridge.events)
|
|
assert gate.active == 1 and gate.waiting == 0
|
|
release.set(); await first.task; await asyncio.sleep(0)
|
|
assert gate.active == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_preparation_exception_releases_reserved_place(monkeypatch):
|
|
gate, bridge, manager, request = await prepare(monkeypatch)
|
|
@bounded_worker_run
|
|
async def fail(*args): raise HTTPException(409, "conflict")
|
|
with pytest.raises(HTTPException): await fail(None, "one", request)
|
|
assert gate.active == gate.waiting == 0
|