mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-22 08:25:21 +08:00
107 lines
3.7 KiB
Python
107 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, Mock
|
|
|
|
import hashlib
|
|
import sqlalchemy
|
|
import pytest
|
|
|
|
from langbot.pkg.api.http.service.apikey import ApiKeyService
|
|
from langbot.pkg.entity.persistence.apikey import ApiKey, ApiKeyStatus
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize('api_key', [None, 123, b'lbk_bytes', '', 'plain_key', ' LBK_bad', 'sk-lbk_fake'])
|
|
async def test_verify_api_key_rejects_non_lbk_keys_without_db_query(api_key):
|
|
persistence_mgr = SimpleNamespace(execute_async=AsyncMock())
|
|
instance_config = SimpleNamespace(data={'api': {'global_api_key': ''}})
|
|
workspace_service = SimpleNamespace(get_execution_binding=AsyncMock())
|
|
service = ApiKeyService(
|
|
SimpleNamespace(
|
|
persistence_mgr=persistence_mgr,
|
|
instance_config=instance_config,
|
|
workspace_service=workspace_service,
|
|
)
|
|
)
|
|
|
|
result = await service.verify_api_key(api_key)
|
|
|
|
assert result is False
|
|
persistence_mgr.execute_async.assert_not_awaited()
|
|
workspace_service.get_execution_binding.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize('key_exists', [True, False])
|
|
async def test_verify_api_key_keeps_db_validation_for_lbk_keys(key_exists):
|
|
key = (
|
|
SimpleNamespace(
|
|
id=1,
|
|
uuid='key-uuid',
|
|
workspace_uuid='workspace-a',
|
|
status=ApiKeyStatus.ACTIVE.value,
|
|
expires_at=None,
|
|
scopes=[],
|
|
)
|
|
if key_exists
|
|
else None
|
|
)
|
|
discovery_result = Mock()
|
|
discovery_result.first.return_value = key
|
|
query_results = [discovery_result]
|
|
if key_exists:
|
|
scoped_result = Mock()
|
|
scoped_result.first.return_value = key
|
|
update_result = Mock()
|
|
update_result.scalar_one_or_none.return_value = key.id
|
|
query_results.extend([scoped_result, update_result])
|
|
persistence_mgr = SimpleNamespace(execute_async=AsyncMock(side_effect=query_results))
|
|
instance_config = SimpleNamespace(data={'api': {'global_api_key': ''}})
|
|
workspace_service = SimpleNamespace(
|
|
get_execution_binding=AsyncMock(
|
|
return_value=SimpleNamespace(
|
|
instance_uuid='instance-a',
|
|
workspace_uuid='workspace-a',
|
|
placement_generation=1,
|
|
)
|
|
)
|
|
)
|
|
service = ApiKeyService(
|
|
SimpleNamespace(
|
|
persistence_mgr=persistence_mgr,
|
|
instance_config=instance_config,
|
|
workspace_service=workspace_service,
|
|
)
|
|
)
|
|
|
|
result = await service.verify_api_key('lbk_valid_format')
|
|
|
|
assert result is key_exists
|
|
digest = hashlib.sha256(b'lbk_valid_format').hexdigest()
|
|
statements = [call.args[0] for call in persistence_mgr.execute_async.await_args_list]
|
|
assert statements[0].whereclause.compare(ApiKey.key_hash == digest)
|
|
if key_exists:
|
|
assert statements[1].whereclause.compare(
|
|
sqlalchemy.and_(
|
|
ApiKey.id == 1,
|
|
ApiKey.workspace_uuid == 'workspace-a',
|
|
ApiKey.key_hash == digest,
|
|
)
|
|
)
|
|
assert statements[2].is_update
|
|
assert statements[2].whereclause.compare(
|
|
sqlalchemy.and_(
|
|
ApiKey.id == 1,
|
|
ApiKey.workspace_uuid == 'workspace-a',
|
|
ApiKey.key_hash == digest,
|
|
ApiKey.status == ApiKeyStatus.ACTIVE.value,
|
|
)
|
|
)
|
|
assert statements[2].compile().params['last_used_at'] is not None
|
|
assert persistence_mgr.execute_async.await_count == (3 if key_exists else 1)
|
|
if key_exists:
|
|
workspace_service.get_execution_binding.assert_awaited_once_with('workspace-a')
|
|
else:
|
|
workspace_service.get_execution_binding.assert_not_awaited()
|