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()