mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +00:00
Hyu
1e6e4c0ca7
* fix(cloud): show owner model balance and enforce single owner * fix(migrations): create owner index idempotently --------- Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
101 lines
4.2 KiB
Python
101 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from langbot.pkg.entity.persistence.base import Base
|
|
from langbot.pkg.entity.persistence.user import User
|
|
from langbot.pkg.entity.persistence.workspace import Workspace, WorkspaceMembership
|
|
from langbot.pkg.persistence.alembic_runner import run_alembic_stamp, run_alembic_upgrade
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_single_owner_migration_demotes_historical_extra_owner_and_installs_unique_index(tmp_path):
|
|
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "single-owner.db"}')
|
|
try:
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
await connection.execute(sa.text('DROP INDEX uq_workspace_memberships_one_active_owner'))
|
|
|
|
session_factory = async_sessionmaker(engine, expire_on_commit=False)
|
|
workspace_uuid = '00000000-0000-4000-8000-000000000001'
|
|
creator_uuid = '00000000-0000-4000-8000-000000000010'
|
|
promoted_uuid = '00000000-0000-4000-8000-000000000020'
|
|
async with session_factory() as session:
|
|
session.add_all(
|
|
[
|
|
User(
|
|
uuid=creator_uuid,
|
|
user='creator@example.test',
|
|
normalized_email='creator@example.test',
|
|
password='hash',
|
|
account_type='local',
|
|
),
|
|
User(
|
|
uuid=promoted_uuid,
|
|
user='promoted@example.test',
|
|
normalized_email='promoted@example.test',
|
|
password='hash',
|
|
account_type='local',
|
|
),
|
|
Workspace(
|
|
uuid=workspace_uuid,
|
|
instance_uuid='instance-test',
|
|
name='Workspace',
|
|
slug='workspace',
|
|
type='team',
|
|
status='active',
|
|
source='local',
|
|
created_by_account_uuid=creator_uuid,
|
|
),
|
|
WorkspaceMembership(
|
|
uuid='00000000-0000-4000-8000-000000000100',
|
|
workspace_uuid=workspace_uuid,
|
|
account_uuid=creator_uuid,
|
|
role='owner',
|
|
status='active',
|
|
),
|
|
WorkspaceMembership(
|
|
uuid='00000000-0000-4000-8000-000000000200',
|
|
workspace_uuid=workspace_uuid,
|
|
account_uuid=promoted_uuid,
|
|
role='owner',
|
|
status='active',
|
|
),
|
|
]
|
|
)
|
|
await session.commit()
|
|
|
|
await run_alembic_stamp(engine, '0018_merge_launch_replay')
|
|
await run_alembic_upgrade(engine, 'head')
|
|
|
|
async with engine.connect() as connection:
|
|
roles = dict(
|
|
(
|
|
await connection.execute(
|
|
sa.text(
|
|
'SELECT account_uuid, role FROM workspace_memberships '
|
|
'WHERE workspace_uuid = :workspace_uuid ORDER BY account_uuid'
|
|
),
|
|
{'workspace_uuid': workspace_uuid},
|
|
)
|
|
).all()
|
|
)
|
|
assert roles == {creator_uuid: 'owner', promoted_uuid: 'admin'}
|
|
indexes = await connection.run_sync(
|
|
lambda sync_connection: {
|
|
index['name'] for index in sa.inspect(sync_connection).get_indexes('workspace_memberships')
|
|
}
|
|
)
|
|
assert 'uq_workspace_memberships_one_active_owner' in indexes
|
|
|
|
with pytest.raises(sa.exc.IntegrityError):
|
|
async with engine.begin() as connection:
|
|
await connection.execute(
|
|
sa.text("UPDATE workspace_memberships SET role = 'owner' WHERE account_uuid = :account_uuid"),
|
|
{'account_uuid': promoted_uuid},
|
|
)
|
|
finally:
|
|
await engine.dispose()
|