mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-22 08:25:21 +08:00
huanghuoguoguo
407e352e5d
test(migrations): reuse historical data contracts and measure startup peaks
279 lines
11 KiB
Python
279 lines
11 KiB
Python
"""
|
|
SQLite migration integration tests.
|
|
|
|
Tests real Alembic migration behavior using temporary SQLite databases.
|
|
Validates the migration workflow from .github/workflows/test-migrations.yml.
|
|
|
|
Run: uv run pytest tests/integration/persistence/test_migrations.py -q
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import pytest
|
|
import sqlalchemy
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
|
|
from langbot.pkg.entity import persistence
|
|
from langbot.pkg.entity.persistence.base import Base
|
|
from langbot.pkg.utils import importutil
|
|
from langbot.pkg.persistence.alembic_runner import (
|
|
run_alembic_downgrade,
|
|
run_alembic_upgrade,
|
|
run_alembic_stamp,
|
|
get_alembic_current,
|
|
_ALEMBIC_DIR,
|
|
)
|
|
from alembic.config import Config
|
|
from alembic.script import ScriptDirectory
|
|
|
|
|
|
# Match PersistenceManager's model registration before create_all, including
|
|
# workspace foreign-key targets, without relying on other tests being collected.
|
|
importutil.import_modules_in_pkg(persistence)
|
|
|
|
|
|
def _get_script_head() -> str:
|
|
"""Resolve the current Alembic head revision from the script directory.
|
|
|
|
Avoids hardcoding a revision number in assertions so adding a new
|
|
migration doesn't require editing the migration tests.
|
|
"""
|
|
cfg = Config()
|
|
cfg.set_main_option('script_location', _ALEMBIC_DIR)
|
|
return ScriptDirectory.from_config(cfg).get_current_head()
|
|
|
|
|
|
pytestmark = pytest.mark.integration
|
|
|
|
|
|
@pytest.fixture
|
|
def sqlite_db_url(tmp_path):
|
|
"""Create SQLite URL with temporary database file."""
|
|
db_file = tmp_path / 'test_migrations.db'
|
|
return f'sqlite+aiosqlite:///{db_file}'
|
|
|
|
|
|
@pytest.fixture
|
|
async def sqlite_engine(sqlite_db_url):
|
|
"""Create async SQLite engine."""
|
|
engine = create_async_engine(sqlite_db_url)
|
|
yield engine
|
|
await engine.dispose()
|
|
|
|
|
|
def test_migration_graph_has_one_head_containing_both_released_branches():
|
|
cfg = Config()
|
|
cfg.set_main_option('script_location', _ALEMBIC_DIR)
|
|
scripts = ScriptDirectory.from_config(cfg)
|
|
heads = scripts.get_heads()
|
|
assert len(heads) == 1, f'Release migrations must converge, found {heads}'
|
|
ancestors = {revision.revision for revision in scripts.walk_revisions()}
|
|
assert {
|
|
'0024_passkey_credentials',
|
|
'0025_bot_plugin_processors',
|
|
'0025_rag_document_identity',
|
|
'0027_pipeline_migration',
|
|
'0028_merge_knowledge_drafts',
|
|
} <= ancestors
|
|
assert all(len(revision) <= 32 for revision in ancestors)
|
|
|
|
|
|
class TestSQLiteMigrationBaseline:
|
|
"""Tests for baseline stamp workflow."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_baseline_stamp_sets_revision(self, sqlite_engine):
|
|
"""
|
|
Stamp baseline on existing tables sets correct revision.
|
|
|
|
Workflow:
|
|
1. Create tables via Base.metadata.create_all
|
|
2. Stamp with '0001_baseline'
|
|
3. Verify current revision is '0001_baseline'
|
|
"""
|
|
# Create all tables (simulates existing DB created by ORM)
|
|
async with sqlite_engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
# Stamp baseline
|
|
await run_alembic_stamp(sqlite_engine, '0001_baseline')
|
|
|
|
# Verify revision
|
|
rev = await get_alembic_current(sqlite_engine)
|
|
assert rev == '0001_baseline', f"Expected '0001_baseline', got {rev}"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_baseline_stamp_on_empty_db(self, sqlite_engine):
|
|
"""
|
|
Stamp on empty database (no tables) still sets revision.
|
|
|
|
This is an edge case - stamping without tables.
|
|
"""
|
|
# Don't create tables - stamp directly
|
|
await run_alembic_stamp(sqlite_engine, '0001_baseline')
|
|
|
|
rev = await get_alembic_current(sqlite_engine)
|
|
assert rev == '0001_baseline'
|
|
|
|
|
|
class TestSQLiteMigrationUpgrade:
|
|
"""Tests for upgrade to head workflow."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upgrade_removes_agent_enabled_column(self, sqlite_engine):
|
|
async with sqlite_engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
await conn.exec_driver_sql('ALTER TABLE agents ADD COLUMN enabled BOOLEAN NOT NULL DEFAULT 1')
|
|
|
|
await run_alembic_stamp(sqlite_engine, '0022_merge_agent_reasoning_heads')
|
|
await run_alembic_upgrade(sqlite_engine, 'head')
|
|
|
|
async with sqlite_engine.connect() as conn:
|
|
columns = await conn.run_sync(
|
|
lambda sync_conn: {column['name'] for column in sqlalchemy.inspect(sync_conn).get_columns('agents')}
|
|
)
|
|
|
|
assert 'enabled' not in columns
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upgrade_from_0012_adds_knowledge_base_embedding_dimension(self, sqlite_engine):
|
|
"""The PostgreSQL pgvector revision also evolves the OSS ORM schema."""
|
|
|
|
async with sqlite_engine.begin() as conn:
|
|
await conn.exec_driver_sql(
|
|
'CREATE TABLE knowledge_bases ('
|
|
'uuid VARCHAR(255) PRIMARY KEY, workspace_uuid VARCHAR(36) NOT NULL, name VARCHAR(255) NOT NULL)'
|
|
)
|
|
|
|
await run_alembic_stamp(sqlite_engine, '0012_plugin_identity')
|
|
await run_alembic_upgrade(sqlite_engine, 'head')
|
|
|
|
async with sqlite_engine.connect() as conn:
|
|
columns = await conn.run_sync(
|
|
lambda sync_conn: {
|
|
item['name'] for item in sqlalchemy.inspect(sync_conn).get_columns('knowledge_bases')
|
|
}
|
|
)
|
|
assert 'embedding_dimension' in columns
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_directory_projection_upgrade_downgrade_round_trip(self, sqlite_engine):
|
|
await run_alembic_stamp(sqlite_engine, '0013_tenant_pgvector')
|
|
|
|
await run_alembic_upgrade(sqlite_engine, 'head')
|
|
async with sqlite_engine.connect() as conn:
|
|
tables = await conn.run_sync(lambda sync_conn: set(sqlalchemy.inspect(sync_conn).get_table_names()))
|
|
assert {'directory_projection_states', 'directory_projection_inbox'} <= tables
|
|
|
|
await run_alembic_downgrade(sqlite_engine, '0013_tenant_pgvector')
|
|
async with sqlite_engine.connect() as conn:
|
|
tables = await conn.run_sync(lambda sync_conn: set(sqlalchemy.inspect(sync_conn).get_table_names()))
|
|
assert 'directory_projection_states' not in tables
|
|
assert 'directory_projection_inbox' not in tables
|
|
|
|
await run_alembic_upgrade(sqlite_engine, 'head')
|
|
assert await get_alembic_current(sqlite_engine) == _get_script_head()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reasoning_config_migrates_existing_models(self, sqlite_engine):
|
|
"""Upgrade from 0017 backfills reasoning config and keeps a database default."""
|
|
async with sqlite_engine.begin() as conn:
|
|
await conn.execute(
|
|
text(
|
|
"""
|
|
CREATE TABLE llm_models (
|
|
uuid VARCHAR(255) PRIMARY KEY,
|
|
name VARCHAR(255) NOT NULL,
|
|
provider_uuid VARCHAR(255) NOT NULL,
|
|
abilities JSON NOT NULL,
|
|
context_length INTEGER,
|
|
extra_args JSON NOT NULL,
|
|
prefered_ranking INTEGER NOT NULL DEFAULT 0
|
|
)
|
|
"""
|
|
)
|
|
)
|
|
await conn.execute(
|
|
text(
|
|
"""
|
|
INSERT INTO llm_models (
|
|
uuid, name, provider_uuid, abilities, extra_args, prefered_ranking
|
|
) VALUES (
|
|
'existing-model', 'Existing Model', 'provider', '[]', '{}', 0
|
|
)
|
|
"""
|
|
)
|
|
)
|
|
|
|
await run_alembic_stamp(sqlite_engine, '0017_oss_workspace_identity')
|
|
await run_alembic_upgrade(sqlite_engine, 'head')
|
|
|
|
async with sqlite_engine.begin() as conn:
|
|
columns = await conn.run_sync(lambda sync_conn: sqlalchemy.inspect(sync_conn).get_columns('llm_models'))
|
|
reasoning_column = next(column for column in columns if column['name'] == 'reasoning_config')
|
|
assert reasoning_column['nullable'] is False
|
|
|
|
existing_value = (
|
|
await conn.execute(text("SELECT reasoning_config FROM llm_models WHERE uuid = 'existing-model'"))
|
|
).scalar_one()
|
|
assert json.loads(existing_value) == {'level': 'provider_default'}
|
|
|
|
await conn.execute(
|
|
text(
|
|
"""
|
|
INSERT INTO llm_models (
|
|
uuid, name, provider_uuid, abilities, extra_args, prefered_ranking
|
|
) VALUES (
|
|
'new-model', 'New Model', 'provider', '[]', '{}', 0
|
|
)
|
|
"""
|
|
)
|
|
)
|
|
new_value = (
|
|
await conn.execute(text("SELECT reasoning_config FROM llm_models WHERE uuid = 'new-model'"))
|
|
).scalar_one()
|
|
assert json.loads(new_value) == {'level': 'provider_default'}
|
|
|
|
|
|
class TestSQLiteMigrationFreshDatabase:
|
|
"""Tests for fresh database workflow."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bot_scoped_sessions_skips_absent_table(self, sqlite_engine):
|
|
"""A partial schema needs no session key migration in either direction."""
|
|
await run_alembic_stamp(sqlite_engine, '0022_codex_credentials')
|
|
await run_alembic_upgrade(sqlite_engine, '0023_bot_scoped_sessions')
|
|
assert await get_alembic_current(sqlite_engine) == '0023_bot_scoped_sessions'
|
|
await run_alembic_downgrade(sqlite_engine, '0022_codex_credentials')
|
|
assert await get_alembic_current(sqlite_engine) == '0022_codex_credentials'
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fresh_db_without_create_all_behavior(self, sqlite_engine):
|
|
"""Bare Alembic reaches head but relies on ORM bootstrap for core tables."""
|
|
await run_alembic_upgrade(sqlite_engine, 'head')
|
|
assert await get_alembic_current(sqlite_engine) == _get_script_head()
|
|
async with sqlite_engine.connect() as conn:
|
|
tables = await conn.run_sync(lambda sync: set(sqlalchemy.inspect(sync).get_table_names()))
|
|
assert {'agents', 'runner_state', 'directory_projection_states', 'directory_projection_inbox'} <= tables
|
|
assert not {'users', 'bots', 'workspaces'} & tables
|
|
|
|
|
|
class TestSQLiteMigrationGetCurrent:
|
|
"""Tests for get_alembic_current behavior."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_current_on_unstamped_db_returns_none(self, sqlite_engine):
|
|
"""
|
|
get_alembic_current returns None for unstamped database.
|
|
"""
|
|
# Create tables but don't stamp
|
|
async with sqlite_engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
# No stamp - should return None
|
|
rev = await get_alembic_current(sqlite_engine)
|
|
assert rev is None, f'Expected None for unstamped DB, got {rev}'
|