LangBot/tests/unit_tests/persistence/test_tenant_uow.py

from __future__ import annotations

import pytest
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine

from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.persistence.tenant_uow import TenantUnitOfWork


pytestmark = pytest.mark.asyncio


async def test_sqlite_tenant_uow_commits_and_rolls_back() -> None:
    engine = create_async_engine('sqlite+aiosqlite:///:memory:')
    table = sa.Table(
        'uow_rows',
        sa.MetaData(),
        sa.Column('id', sa.Integer, primary_key=True),
        sa.Column('workspace_uuid', sa.String(36), nullable=False),
    )
    try:
        async with engine.begin() as conn:
            await conn.run_sync(table.metadata.create_all)

        async with TenantUnitOfWork(engine, 'workspace-a') as uow:
            await uow.execute(sa.insert(table).values(id=1, workspace_uuid='workspace-a'))

        with pytest.raises(RuntimeError, match='roll back'):
            async with TenantUnitOfWork(engine, 'workspace-a') as uow:
                await uow.execute(sa.insert(table).values(id=2, workspace_uuid='workspace-a'))
                raise RuntimeError('roll back this transaction')

        async with engine.connect() as conn:
            rows = (await conn.execute(sa.select(table.c.id).order_by(table.c.id))).scalars().all()
        assert rows == [1]
    finally:
        await engine.dispose()


async def test_tenant_uow_is_single_use_and_requires_an_active_scope() -> None:
    engine = create_async_engine('sqlite+aiosqlite:///:memory:')
    uow = TenantUnitOfWork(engine, 'workspace-a')
    try:
        with pytest.raises(RuntimeError, match='not active'):
            _ = uow.session

        async with uow:
            assert uow.session.in_transaction()

        with pytest.raises(RuntimeError, match='cannot be reused'):
            async with uow:
                pass
    finally:
        await engine.dispose()


def test_persistence_mode_must_be_a_trusted_enum() -> None:
    with pytest.raises(TypeError, match='trusted PersistenceMode'):
        PersistenceManager(object(), mode='cloud_runtime')  # type: ignore[arg-type]

    manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
    assert manager.mode is PersistenceMode.CLOUD_RUNTIME