LangBot/tests/unit_tests/api/test_apikey_service.py

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock

import pytest

from langbot.pkg.api.http.service.apikey import ApiKeyService
from langbot.pkg.entity.persistence.apikey import 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
    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()