mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 04:40:57 +00:00
dadachann
103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, Mock
|
|
|
|
import pytest
|
|
import sqlalchemy
|
|
|
|
from langbot.pkg.cloud.entitlements import EntitlementResolver, EntitlementSnapshot
|
|
from langbot.pkg.cloud.quotas import (
|
|
WorkspaceQuota,
|
|
require_resource_capacity,
|
|
resolve_workspace_quota,
|
|
)
|
|
from langbot.pkg.entity.persistence.bot import Bot
|
|
|
|
|
|
WORKSPACE_UUID = '11111111-1111-1111-1111-111111111111'
|
|
INSTANCE_UUID = 'cloud-instance'
|
|
|
|
|
|
class _Provider:
|
|
def __init__(self, limits: dict[str, int]) -> None:
|
|
self.limits = limits
|
|
|
|
async def get_workspace_entitlement(self, workspace_uuid: str) -> EntitlementSnapshot:
|
|
return EntitlementSnapshot(
|
|
instance_uuid=INSTANCE_UUID,
|
|
workspace_uuid=workspace_uuid,
|
|
entitlement_revision=1,
|
|
status='active',
|
|
not_before=0,
|
|
expires_at=4_102_444_800,
|
|
features={},
|
|
limits=self.limits,
|
|
plan_name='test',
|
|
)
|
|
|
|
|
|
async def _resolver(limits: dict[str, int]) -> EntitlementResolver:
|
|
resolver = EntitlementResolver(INSTANCE_UUID, _Provider(limits))
|
|
await resolver.reconcile_active_workspaces({WORKSPACE_UUID})
|
|
return resolver
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_workspace_quota_uses_signed_cloud_limit() -> None:
|
|
ap = SimpleNamespace(entitlement_resolver=await _resolver({'bots.max': 2}))
|
|
|
|
quota = await resolve_workspace_quota(ap, WORKSPACE_UUID, 'bots.max', fallback=99)
|
|
|
|
assert quota == WorkspaceQuota(limit=2, requires_transaction_lock=True)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_workspace_quota_preserves_oss_fallback() -> None:
|
|
quota = await resolve_workspace_quota(SimpleNamespace(), WORKSPACE_UUID, 'bots.max', fallback=7)
|
|
|
|
assert quota == WorkspaceQuota(limit=7, requires_transaction_lock=False)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_require_resource_capacity_locks_workspace_before_counting() -> None:
|
|
statements: list[object] = []
|
|
lock_result = Mock()
|
|
lock_result.first.return_value = (WORKSPACE_UUID,)
|
|
count_result = Mock()
|
|
count_result.scalar_one.return_value = 1
|
|
execute = AsyncMock(side_effect=[lock_result, count_result])
|
|
|
|
await require_resource_capacity(
|
|
execute,
|
|
workspace_uuid=WORKSPACE_UUID,
|
|
model=Bot,
|
|
quota=WorkspaceQuota(limit=2, requires_transaction_lock=True),
|
|
resource_name='bots',
|
|
)
|
|
|
|
statements.extend(call.args[0] for call in execute.await_args_list)
|
|
assert len(statements) == 2
|
|
assert isinstance(statements[0], sqlalchemy.sql.Select)
|
|
assert statements[0]._for_update_arg is not None
|
|
assert 'workspaces' in str(statements[0])
|
|
assert 'count' in str(statements[1]).lower()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_require_resource_capacity_rejects_at_boundary() -> None:
|
|
lock_result = Mock()
|
|
lock_result.first.return_value = (WORKSPACE_UUID,)
|
|
count_result = Mock()
|
|
count_result.scalar_one.return_value = 2
|
|
execute = AsyncMock(side_effect=[lock_result, count_result])
|
|
|
|
with pytest.raises(ValueError, match=r'Maximum number of bots \(2\) reached'):
|
|
await require_resource_capacity(
|
|
execute,
|
|
workspace_uuid=WORKSPACE_UUID,
|
|
model=Bot,
|
|
quota=WorkspaceQuota(limit=2, requires_transaction_lock=True),
|
|
resource_name='bots',
|
|
)
|