vent_agent_2.0/tests/test_environment.py

"""公共生命周期：启动失败、取消、重复进入以及真实 SQLite 资源释放。"""
import asyncio
from concurrent.futures import ThreadPoolExecutor
import threading
from types import SimpleNamespace

import pytest
from langchain_core.messages import AIMessageChunk

import config
from agents import agent_memory, agent_registry
from db import db_conn, chat_store, billing_store
from runtime import environment
from runtime.environment import RuntimeEnvironment
from runtime.execution import execute
from scheduler import scheduler


@pytest.fixture
def lifecycle(monkeypatch):
    calls = []
    def mark(name):
        def call(*args, **kwargs):
            calls.append(name)
        return call
    def amark(name):
        async def call(*args, **kwargs):
            calls.append(name)
        return call
    for name in ('setup_logging', 'load_extensions', 'approval_names'):
        monkeypatch.setattr(environment, name, mark(name))
    monkeypatch.setattr(config, 'log_config_sources', mark('config'))
    monkeypatch.setattr(db_conn, 'init_schema', mark('schema'))
    monkeypatch.setattr(db_conn, 'health', lambda: {'status': 'ok'})
    monkeypatch.setattr(db_conn, 'close', mark('db_close'))
    monkeypatch.setattr(environment.pricing, 'seed_default_pricing', mark('pricing'))
    monkeypatch.setattr(environment.run_store, 'recover_interrupted', mark('recover'))
    monkeypatch.setattr(agent_registry, 'invalidate', mark('invalidate'))
    monkeypatch.setattr(agent_memory, 'init_checkpointer', amark('checkpoint'))
    monkeypatch.setattr(agent_memory, 'close_checkpointer', amark('checkpoint_close'))
    monkeypatch.setattr(scheduler, 'start', amark('scheduler_start'))
    monkeypatch.setattr(scheduler, 'shutdown', amark('scheduler_stop'))
    return calls


@pytest.mark.asyncio
async def test_default_environment_and_sequential_reuse(lifecycle):
    env = RuntimeEnvironment(prewarm=False, billing_retry=False)
    for _ in range(2):
        lifecycle.clear()
        async with env as current:
            assert current is env
            assert lifecycle[-4:] == ['pricing', 'recover', 'invalidate', 'checkpoint']
        assert lifecycle[-3:] == ['invalidate', 'checkpoint_close', 'db_close']
        assert 'scheduler_start' not in lifecycle


@pytest.mark.asyncio
async def test_nested_environment_does_not_close_outer(lifecycle):
    env = RuntimeEnvironment(prewarm=False, billing_retry=False)
    async with env:
        for other in (env, RuntimeEnvironment()):
            with pytest.raises(RuntimeError, match='已有活动'):
                async with other:
                    pytest.fail('不应启动嵌套环境')
        assert 'checkpoint_close' not in lifecycle
    assert lifecycle.count('db_close') == 1


@pytest.mark.asyncio
@pytest.mark.parametrize('stage', ['pricing', 'checkpoint', 'scheduler'])
async def test_partial_startup_failure_cleans_and_allows_restart(stage, lifecycle, monkeypatch):
    def fail():
        raise ValueError(stage)
    async def afail():
        raise ValueError(stage)
    with monkeypatch.context() as patch:
        if stage == 'pricing':
            patch.setattr(environment.pricing, 'seed_default_pricing', fail)
        elif stage == 'checkpoint':
            patch.setattr(agent_memory, 'init_checkpointer', afail)
        else:
            patch.setattr(scheduler, 'start', afail)
        with pytest.raises(ValueError, match=stage):
            async with RuntimeEnvironment(scheduler=True, prewarm=False, billing_retry=False):
                pytest.fail('启动失败后不应进入执行体')
    assert lifecycle[-1] == 'db_close'
    if stage != 'pricing':
        assert 'checkpoint_close' in lifecycle
    if stage == 'scheduler':
        assert lifecycle.index('scheduler_stop') < lifecycle.index('checkpoint_close')
    async with RuntimeEnvironment(prewarm=False, billing_retry=False):
        pass


@pytest.mark.asyncio
async def test_body_exception_and_cleanup_error_still_release_resources(lifecycle, monkeypatch):
    async def fail_close():
        lifecycle.append('checkpoint_close')
        raise RuntimeError('close failed')
    with monkeypatch.context() as patch:
        patch.setattr(agent_memory, 'close_checkpointer', fail_close)
        with pytest.raises(RuntimeError, match='close failed'):
            async with RuntimeEnvironment(prewarm=False, billing_retry=False):
                raise ValueError('body failed')
    assert lifecycle[-1] == 'db_close'
    async with RuntimeEnvironment(prewarm=False, billing_retry=False):
        pass


@pytest.mark.asyncio
async def test_cancel_during_startup_waits_for_initialization(lifecycle, monkeypatch):
    started, finish = asyncio.Event(), asyncio.Event()
    async def init():
        started.set()
        await finish.wait()
        lifecycle.append('checkpoint')
    monkeypatch.setattr(agent_memory, 'init_checkpointer', init)
    async def run():
        async with RuntimeEnvironment(prewarm=False, billing_retry=False):
            pytest.fail('取消后不应进入执行体')
    task = asyncio.create_task(run())
    await started.wait()
    task.cancel()
    await asyncio.sleep(0)
    task.cancel()  # 清理期间再次取消，也要等待底层初始化收尾。
    assert 'db_close' not in lifecycle
    finish.set()
    with pytest.raises(asyncio.CancelledError):
        await task
    assert lifecycle[-4:] == ['checkpoint', 'invalidate', 'checkpoint_close', 'db_close']


@pytest.mark.asyncio
async def test_shutdown_waits_for_workers_before_closing_connections(lifecycle, monkeypatch):
    from billing import meter
    prewarm_started, billing_started = threading.Event(), threading.Event()
    release = threading.Event()
    def compile_agent():
        prewarm_started.set()
        assert release.wait(5)
        lifecycle.append('compiled')
    def retry():
        billing_started.set()
        assert release.wait(5)
        lifecycle.append('settled')
    monkeypatch.setattr(agent_registry, 'get_agent', compile_agent)
    monkeypatch.setattr(meter, 'retry_pending', retry)
    running = asyncio.Event()
    async def run():
        async with RuntimeEnvironment(scheduler=True):
            running.set()
            await asyncio.Event().wait()
    task = asyncio.create_task(run())
    try:
        await running.wait()
        assert await asyncio.to_thread(prewarm_started.wait, 2)
        assert await asyncio.to_thread(billing_started.wait, 2)
        task.cancel()
        await asyncio.sleep(0)
        await asyncio.sleep(0)
        assert 'checkpoint_close' not in lifecycle
    finally:
        release.set()
    with pytest.raises(asyncio.CancelledError):
        await task
    assert lifecycle.index('scheduler_stop') < lifecycle.index('checkpoint_close')
    assert lifecycle.index('settled') < lifecycle.index('db_close')
    assert lifecycle.index('compiled') < lifecycle.index('checkpoint_close')
    assert not [t for t in asyncio.all_tasks() if t.get_name().startswith('runtime-')]


class FakeAgent:
    def __init__(self):
        self.checkpointer = agent_memory.get_checkpointer()

    async def aget_state(self, options):
        return SimpleNamespace(tasks=[])

    async def astream(self, *args, **kwargs):
        yield {'type': 'messages', 'data': (AIMessageChunk(content='报告完成', usage_metadata={
            'input_tokens': 8, 'output_tokens': 2, 'total_tokens': 10}), {})}


@pytest.mark.asyncio
async def test_sqlite_script_executes_bills_and_reopens(monkeypatch, tmp_path):
    # 独立业务库与 checkpoint，不使用线上配置，也不调用模型网络。
    adapter = db_conn._SqliteAdapter(str(tmp_path / 'business.sqlite'))
    monkeypatch.setattr(db_conn, '_adapter', adapter)
    monkeypatch.setattr(db_conn, '_schema_ready', False)
    monkeypatch.setattr(agent_memory, '_saver', None)
    monkeypatch.setattr(agent_memory, '_kind', 'sqlite')
    monkeypatch.setattr(config, 'CHECKPOINT_DB_PATH', str(tmp_path / 'checkpoint.sqlite'))
    monkeypatch.setattr(agent_registry._slot, '_compile', FakeAgent)
    schema = config.ROOT / 'db' / 'schema_sqlite.sql'
    monkeypatch.setattr(adapter, 'schema_file', lambda: schema)
    monkeypatch.setattr(config, 'ROOT', tmp_path)
    previous = None
    env = RuntimeEnvironment(prewarm=False, billing_retry=False)
    for _ in range(2):
        async with env:
            session = chat_store.create_session('runtime-test')
            agent = agent_registry.get_agent()
            assert agent is not previous
            assert agent.checkpointer is agent_memory.get_checkpointer()
            result = await execute(user_name='runtime-test', session_id=session['id'], message='生成报告')
            assert result.status == 'success'
            assert result.output == '报告完成'
            assert result.usage['total_tokens'] == 10
            assert billing_store.list_pricing()
            assert db_conn.scalar('SELECT total_tokens FROM billing_usage WHERE session_id = ?',
                                  (session['id'],)) == 10
            conn = adapter._conn
            saver_conn = agent.checkpointer.conn
            previous = agent
        assert adapter._conn is None
        with pytest.raises(Exception, match='closed'):
            conn.execute('SELECT 1')
        with pytest.raises(ValueError, match='no active connection'):
            await saver_conn.execute('SELECT 1')
        assert not agent_memory.is_initialized()
        assert agent_registry.peek_agent() is None


def test_mysql_closes_connections_from_all_threads_and_reopens():
    adapter = db_conn._MySqlAdapter()
    class Connection:
        def __init__(self):
            self.closed = False
        def close(self):
            self.closed = True
        def ping(self, **kwargs):
            assert not self.closed
    adapter._connect = lambda database: Connection()
    barrier = threading.Barrier(3)
    def connect():
        conn = adapter._connection()
        barrier.wait(timeout=3)
        return conn
    with ThreadPoolExecutor(max_workers=2) as pool:
        tasks = [pool.submit(connect) for _ in range(2)]
        main_conn = adapter._connection()
        barrier.wait(timeout=3)
        connections = [main_conn, *(task.result() for task in tasks)]
        assert len({id(conn) for conn in connections}) == 3
        adapter.close()
        assert all(conn.closed for conn in connections)
        assert pool.submit(adapter._connection).result() not in connections
        assert adapter._connection() not in connections
        adapter.close()


@pytest.mark.asyncio
async def test_checkpoint_setup_failure_closes_partial_connection(monkeypatch, tmp_path):
    import aiosqlite
    from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
    monkeypatch.setattr(agent_memory, '_saver', None)
    monkeypatch.setattr(agent_memory, '_kind', 'sqlite')
    monkeypatch.setattr(config, 'CHECKPOINT_DB_PATH', str(tmp_path / 'bad.sqlite'))
    connected = []
    original = aiosqlite.connect
    async def connect(*args, **kwargs):
        conn = await original(*args, **kwargs)
        connected.append(conn)
        return conn
    async def fail(self):
        raise ValueError('setup failed')
    monkeypatch.setattr(aiosqlite, 'connect', connect)
    monkeypatch.setattr(AsyncSqliteSaver, 'setup', fail)
    with pytest.raises(agent_memory.CheckpointerConfigError, match='setup failed'):
        await agent_memory.init_checkpointer()
    assert not agent_memory.is_initialized()
    with pytest.raises(ValueError, match='no active connection'):
        await connected[0].execute('SELECT 1')


@pytest.mark.asyncio
async def test_scheduler_start_cancellation_waits_for_database_read(lifecycle, monkeypatch):
    started, finish = asyncio.Event(), asyncio.Event()
    async def start():
        started.set()
        await finish.wait()
        lifecycle.append('scheduler_loaded')
    monkeypatch.setattr(scheduler, 'start', start)
    async def run():
        async with RuntimeEnvironment(scheduler=True, prewarm=False, billing_retry=False):
            pytest.fail('取消后不应进入执行体')
    task = asyncio.create_task(run())
    await started.wait()
    task.cancel()
    await asyncio.sleep(0)
    assert 'db_close' not in lifecycle
    finish.set()
    with pytest.raises(asyncio.CancelledError):
        await task
    assert lifecycle.index('scheduler_loaded') < lifecycle.index('scheduler_stop')
    assert lifecycle.index('scheduler_stop') < lifecycle.index('db_close')


@pytest.mark.asyncio
async def test_http_lifespan_uses_common_environment(lifecycle, monkeypatch):
    import main
    def env(**kwargs):
        assert kwargs == {'scheduler': True}
        return RuntimeEnvironment(**kwargs, prewarm=False, billing_retry=False)
    monkeypatch.setattr(main, 'RuntimeEnvironment', env)
    async with main.app.router.lifespan_context(main.app):
        assert 'scheduler_start' in lifecycle
        assert 'db_close' not in lifecycle
    assert lifecycle[-4:] == ['scheduler_stop', 'invalidate', 'checkpoint_close', 'db_close']


def test_environment_reopens_across_event_loops(lifecycle):
    env = RuntimeEnvironment(prewarm=False, billing_retry=False, recover_interrupted=False)
    async def run():
        async with env:
            pass
    asyncio.run(run())
    asyncio.run(run())
    assert lifecycle.count('checkpoint') == 2
    assert lifecycle.count('db_close') == 2
    assert 'recover' not in lifecycle