mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-17 03:04:20 +00:00
huanghuoguoguo
246 lines
8.3 KiB
Python
246 lines
8.3 KiB
Python
"""Tests for RunLedgerStore host primitives."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import datetime
|
|
|
|
import pytest
|
|
import sqlalchemy
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from langbot.pkg.agent.runner.run_ledger_store import RunLedgerStore
|
|
from langbot.pkg.entity.persistence.agent_run import AgentRun
|
|
from langbot.pkg.entity.persistence.base import Base
|
|
|
|
|
|
UTC = datetime.timezone.utc
|
|
|
|
|
|
@pytest.fixture
|
|
async def db_engine(tmp_path):
|
|
db_path = tmp_path / 'run_ledger_store.db'
|
|
engine = create_async_engine(f'sqlite+aiosqlite:///{db_path}', echo=False)
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
yield engine
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.fixture
|
|
def store(db_engine):
|
|
return RunLedgerStore(db_engine)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_queued_run_claim_renew_release(store):
|
|
run = await store.create_run(
|
|
run_id='run-queued',
|
|
event_id='evt-1',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
status='queued',
|
|
queue_name='default',
|
|
priority=10,
|
|
requested_runtime_id='runtime-a',
|
|
)
|
|
|
|
assert run['status'] == 'queued'
|
|
assert run['started_at'] is None
|
|
assert run['queue_name'] == 'default'
|
|
assert run['priority'] == 10
|
|
assert run['requested_runtime_id'] == 'runtime-a'
|
|
|
|
assert await store.claim_next_run(runtime_id='runtime-b', queue_name='default') is None
|
|
|
|
claimed = await store.claim_next_run(runtime_id='runtime-a', queue_name='default', lease_seconds=30)
|
|
|
|
assert claimed is not None
|
|
assert claimed['run_id'] == 'run-queued'
|
|
assert claimed['status'] == 'claimed'
|
|
assert claimed['claimed_by_runtime_id'] == 'runtime-a'
|
|
assert claimed['claim_token']
|
|
assert claimed['dispatch_attempts'] == 1
|
|
assert claimed['claim_lease_expires_at'] is not None
|
|
assert claimed['last_claimed_at'] is not None
|
|
|
|
token = claimed['claim_token']
|
|
assert await store.renew_claim(run_id='run-queued', claim_token='wrong-token') is None
|
|
|
|
renewed = await store.renew_claim(run_id='run-queued', claim_token=token, lease_seconds=90)
|
|
|
|
assert renewed is not None
|
|
assert renewed['claim_token'] == token
|
|
assert renewed['claim_lease_expires_at'] >= claimed['claim_lease_expires_at']
|
|
|
|
released = await store.release_claim(
|
|
run_id='run-queued',
|
|
claim_token=token,
|
|
status='queued',
|
|
status_reason='runtime released capacity',
|
|
)
|
|
|
|
assert released is not None
|
|
assert released['status'] == 'queued'
|
|
assert released['status_reason'] == 'runtime released capacity'
|
|
assert released['claimed_by_runtime_id'] is None
|
|
assert released['claim_token'] is None
|
|
assert released['claim_lease_expires_at'] is None
|
|
assert released['dispatch_attempts'] == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_expired_claim_can_be_reclaimed(store, db_engine):
|
|
await store.create_run(
|
|
run_id='run-expired',
|
|
event_id='evt-2',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
status='queued',
|
|
queue_name='default',
|
|
)
|
|
first_claim = await store.claim_next_run(runtime_id='runtime-a', queue_name='default', lease_seconds=60)
|
|
assert first_claim is not None
|
|
|
|
session_factory = sessionmaker(db_engine, class_=AsyncSession, expire_on_commit=False)
|
|
async with session_factory() as session:
|
|
await session.execute(
|
|
sqlalchemy.update(AgentRun)
|
|
.where(AgentRun.run_id == 'run-expired')
|
|
.values(claim_lease_expires_at=datetime.datetime.now(UTC) - datetime.timedelta(seconds=1))
|
|
)
|
|
await session.commit()
|
|
|
|
reclaimed = await store.claim_next_run(runtime_id='runtime-b', queue_name='default', lease_seconds=60)
|
|
|
|
assert reclaimed is not None
|
|
assert reclaimed['run_id'] == 'run-expired'
|
|
assert reclaimed['claimed_by_runtime_id'] == 'runtime-b'
|
|
assert reclaimed['claim_token'] != first_claim['claim_token']
|
|
assert reclaimed['dispatch_attempts'] == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_release_expired_claims_requeues_runs(store, db_engine):
|
|
await store.create_run(
|
|
run_id='run-expired-release',
|
|
event_id='evt-3',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
status='queued',
|
|
queue_name='default',
|
|
)
|
|
await store.create_run(
|
|
run_id='run-active-claim',
|
|
event_id='evt-4',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
status='queued',
|
|
queue_name='default',
|
|
)
|
|
expired_claim = await store.claim_next_run(runtime_id='runtime-a', queue_name='default', lease_seconds=60)
|
|
active_claim = await store.claim_next_run(runtime_id='runtime-b', queue_name='default', lease_seconds=60)
|
|
assert expired_claim is not None
|
|
assert active_claim is not None
|
|
|
|
session_factory = sessionmaker(db_engine, class_=AsyncSession, expire_on_commit=False)
|
|
async with session_factory() as session:
|
|
await session.execute(
|
|
sqlalchemy.update(AgentRun)
|
|
.where(AgentRun.run_id == 'run-expired-release')
|
|
.values(claim_lease_expires_at=datetime.datetime.now(UTC) - datetime.timedelta(seconds=1))
|
|
)
|
|
await session.commit()
|
|
|
|
released = await store.release_expired_claims()
|
|
|
|
assert [run['run_id'] for run in released] == ['run-expired-release']
|
|
assert released[0]['status'] == 'queued'
|
|
assert released[0]['status_reason'] == 'claim lease expired'
|
|
assert released[0]['claimed_by_runtime_id'] is None
|
|
assert released[0]['claim_token'] is None
|
|
assert released[0]['claim_lease_expires_at'] is None
|
|
|
|
active = await store.get_run('run-active-claim')
|
|
assert active is not None
|
|
assert active['status'] == 'claimed'
|
|
assert active['claim_token'] == active_claim['claim_token']
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_append_audit_event_uses_next_sequence(store):
|
|
await store.create_run(
|
|
run_id='run-audit',
|
|
event_id='evt-5',
|
|
binding_id='binding-1',
|
|
runner_id='runner-a',
|
|
)
|
|
await store.append_event(
|
|
run_id='run-audit',
|
|
sequence=1,
|
|
event_type='message.completed',
|
|
data={'ok': True},
|
|
)
|
|
|
|
event = await store.append_audit_event(
|
|
run_id='run-audit',
|
|
event_type='admin.run_cancel',
|
|
data={'action': 'run_cancel'},
|
|
metadata={'permission': 'agent_run:admin'},
|
|
)
|
|
|
|
assert event is not None
|
|
assert event['sequence'] == 2
|
|
assert event['type'] == 'admin.run_cancel'
|
|
assert event['source'] == 'host'
|
|
assert event['data'] == {'action': 'run_cancel'}
|
|
assert event['metadata'] == {'permission': 'agent_run:admin'}
|
|
assert await store.append_audit_event(run_id='missing', event_type='admin.missing') is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_runtime_register_heartbeat_list_and_mark_stale(store):
|
|
registered = await store.register_runtime(
|
|
runtime_id='runtime-a',
|
|
display_name='Runtime A',
|
|
endpoint='http://runtime-a',
|
|
version='1.0.0',
|
|
capabilities={'stream': True},
|
|
labels={'region': 'test'},
|
|
metadata={'slot_count': 2},
|
|
heartbeat_deadline_seconds=30,
|
|
)
|
|
|
|
assert registered['runtime_id'] == 'runtime-a'
|
|
assert registered['status'] == 'online'
|
|
assert registered['display_name'] == 'Runtime A'
|
|
assert registered['capabilities'] == {'stream': True}
|
|
assert registered['labels'] == {'region': 'test'}
|
|
assert registered['metadata'] == {'slot_count': 2}
|
|
assert registered['last_heartbeat_at'] is not None
|
|
assert registered['heartbeat_deadline_at'] is not None
|
|
|
|
heartbeat = await store.heartbeat_runtime(
|
|
runtime_id='runtime-a',
|
|
metadata={'active_runs': 1},
|
|
heartbeat_deadline_seconds=30,
|
|
)
|
|
|
|
assert heartbeat is not None
|
|
assert heartbeat['metadata'] == {'slot_count': 2, 'active_runs': 1}
|
|
|
|
runtimes = await store.list_runtimes(statuses=['online'])
|
|
assert [runtime['runtime_id'] for runtime in runtimes] == ['runtime-a']
|
|
|
|
stale = await store.mark_stale_runtimes(
|
|
now=datetime.datetime.now(UTC) + datetime.timedelta(seconds=31),
|
|
)
|
|
|
|
assert [runtime['runtime_id'] for runtime in stale] == ['runtime-a']
|
|
assert stale[0]['status'] == 'stale'
|
|
assert (await store.get_runtime('runtime-a'))['status'] == 'stale'
|