""" Unit tests for LiteLlmManager class. """ import importlib import os import sys from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from pydantic import SecretStr from server.constants import ( get_default_litellm_model, ) from storage.lite_llm_manager import ( LiteLlmManager, get_byor_key_alias, get_openhands_cloud_key_alias, ) from storage.user_settings import UserSettings from openhands.app_server.settings.settings_models import Settings def _agent_value(settings: Settings, key: str): """Navigate into settings.agent_settings using a dot-separated key.""" obj = settings.agent_settings for part in key.split('.'): obj = getattr(obj, part) return obj def _secret_value(settings: Settings, key: str): """Navigate into settings.agent_settings and unwrap SecretStr values.""" secret = _agent_value(settings, key) return secret.get_secret_value() if secret else None class TestDefaultInitialBudget: """Test cases for DEFAULT_INITIAL_BUDGET configuration.""" @pytest.fixture(autouse=True) def restore_module_state(self): """Ensure module is properly restored after each test.""" # Save original module if it exists original_module = sys.modules.get('storage.lite_llm_manager') yield # Restore module state after each test if 'storage.lite_llm_manager' in sys.modules: del sys.modules['storage.lite_llm_manager'] # Clear the env vars os.environ.pop('DEFAULT_INITIAL_BUDGET', None) os.environ.pop('ENABLE_BILLING', None) # Restore original module or reimport fresh if original_module is not None: sys.modules['storage.lite_llm_manager'] = original_module else: importlib.import_module('storage.lite_llm_manager') def test_default_initial_budget_none_when_billing_disabled(self): """Test that DEFAULT_INITIAL_BUDGET is None when billing is disabled.""" # Temporarily remove the module so we can reimport with different env vars if 'storage.lite_llm_manager' in sys.modules: del sys.modules['storage.lite_llm_manager'] # Ensure billing is disabled (default) and reimport os.environ.pop('ENABLE_BILLING', None) os.environ.pop('DEFAULT_INITIAL_BUDGET', None) module = importlib.import_module('storage.lite_llm_manager') assert module.DEFAULT_INITIAL_BUDGET is None def test_default_initial_budget_defaults_to_zero_when_billing_enabled(self): """Test that DEFAULT_INITIAL_BUDGET defaults to 0.0 when billing is enabled.""" # Temporarily remove the module so we can reimport with different env vars if 'storage.lite_llm_manager' in sys.modules: del sys.modules['storage.lite_llm_manager'] # Enable billing and reimport os.environ['ENABLE_BILLING'] = 'true' os.environ.pop('DEFAULT_INITIAL_BUDGET', None) module = importlib.import_module('storage.lite_llm_manager') assert module.DEFAULT_INITIAL_BUDGET == 0.0 def test_default_initial_budget_uses_env_var_when_billing_enabled(self): """Test that DEFAULT_INITIAL_BUDGET uses value from environment variable when billing enabled.""" if 'storage.lite_llm_manager' in sys.modules: del sys.modules['storage.lite_llm_manager'] os.environ['ENABLE_BILLING'] = 'true' os.environ['DEFAULT_INITIAL_BUDGET'] = '100.0' module = importlib.import_module('storage.lite_llm_manager') assert module.DEFAULT_INITIAL_BUDGET == 100.0 def test_default_initial_budget_ignores_env_var_when_billing_disabled(self): """Test that DEFAULT_INITIAL_BUDGET returns None when billing disabled, ignoring env var.""" if 'storage.lite_llm_manager' in sys.modules: del sys.modules['storage.lite_llm_manager'] os.environ.pop('ENABLE_BILLING', None) # billing disabled by default os.environ['DEFAULT_INITIAL_BUDGET'] = '100.0' module = importlib.import_module('storage.lite_llm_manager') assert module.DEFAULT_INITIAL_BUDGET is None def test_default_initial_budget_rejects_invalid_value(self): """Test that DEFAULT_INITIAL_BUDGET raises ValueError for invalid values.""" if 'storage.lite_llm_manager' in sys.modules: del sys.modules['storage.lite_llm_manager'] os.environ['ENABLE_BILLING'] = 'true' os.environ['DEFAULT_INITIAL_BUDGET'] = 'abc' with pytest.raises(ValueError) as exc_info: importlib.import_module('storage.lite_llm_manager') assert 'Invalid DEFAULT_INITIAL_BUDGET' in str(exc_info.value) def test_default_initial_budget_rejects_negative_value(self): """Test that DEFAULT_INITIAL_BUDGET raises ValueError for negative values.""" if 'storage.lite_llm_manager' in sys.modules: del sys.modules['storage.lite_llm_manager'] os.environ['ENABLE_BILLING'] = 'true' os.environ['DEFAULT_INITIAL_BUDGET'] = '-10.0' with pytest.raises(ValueError) as exc_info: importlib.import_module('storage.lite_llm_manager') assert 'must be non-negative' in str(exc_info.value) class TestLiteLlmManager: """Test cases for LiteLlmManager class.""" @pytest.fixture def mock_settings(self): """Create a mock Settings object.""" settings = Settings() settings.update( { 'agent_settings_diff': { 'agent': 'TestAgent', 'llm': { 'model': 'test-model', 'api_key': 'test-key', 'base_url': 'http://test.com', }, }, } ) return settings @pytest.fixture def mock_user_settings(self): """Create a mock UserSettings object.""" user_settings = UserSettings() user_settings.agent_settings = { 'agent': 'TestAgent', 'llm': { 'model': 'test-model', 'base_url': 'http://test.com', 'api_key': 'test-key', }, } user_settings.llm_api_key = SecretStr('test-key') user_settings.user_version = 4 # Set version to avoid None comparison return user_settings @pytest.fixture def mock_http_client(self): """Create a mock HTTP client.""" client = AsyncMock(spec=httpx.AsyncClient) return client @pytest.fixture def mock_response(self): """Create a mock HTTP response.""" response = MagicMock() response.is_success = True response.status_code = 200 response.text = 'Success' response.json.return_value = {'key': 'test-api-key'} response.raise_for_status = MagicMock() return response @pytest.fixture def mock_team_response(self): """Create a mock team response.""" response = MagicMock() response.is_success = True response.status_code = 200 response.json.return_value = { 'team_memberships': [ { 'user_id': 'test-user-id', 'team_id': 'test-org-id', 'max_budget': 100.0, } ] } response.raise_for_status = MagicMock() return response @pytest.fixture def mock_user_response(self): """Create a mock user response.""" response = MagicMock() response.is_success = True response.status_code = 200 response.json.return_value = { 'user_info': { 'max_budget': 50.0, 'spend': 10.0, } } response.raise_for_status = MagicMock() return response @pytest.fixture def mock_key_info_response(self): """Create a mock key info response.""" response = MagicMock() response.is_success = True response.status_code = 200 response.json.return_value = { 'info': { 'max_budget': 100.0, 'spend': 25.0, } } response.raise_for_status = MagicMock() return response @pytest.mark.asyncio async def test_create_entries_missing_config(self, mock_settings): """Test create_entries when LiteLLM config is missing.""" with patch.dict(os.environ, {'LITE_LLM_API_KEY': '', 'LITE_LLM_API_URL': ''}): with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', None): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', None): result = await LiteLlmManager.create_entries( 'test-org-id', 'test-user-id', mock_settings, create_user=True ) assert result is None @pytest.mark.asyncio async def test_create_entries_direct_defaults_skip_litellm(self, mock_settings): """Test direct LLM defaults without provisioning LiteLLM users or keys.""" with ( patch( 'storage.lite_llm_manager.should_use_direct_llm_defaults', return_value=True, ), patch( 'storage.lite_llm_manager.get_default_llm_model', return_value='openai/anthropic/claude-sonnet', ), patch( 'storage.lite_llm_manager.get_default_llm_base_url', return_value='https://bifrost.example.com/openai/v1', ), patch( 'storage.lite_llm_manager.get_default_llm_api_key', return_value='sk-bf-shared-smoke-test', ), patch('httpx.AsyncClient') as mock_client_class, ): result = await LiteLlmManager.create_entries( 'test-org-id', 'test-user-id', mock_settings, create_user=True ) assert result is not None assert _agent_value(result, 'agent') == 'CodeActAgent' assert _agent_value(result, 'llm.model') == 'openai/anthropic/claude-sonnet' assert ( _agent_value(result, 'llm.base_url') == 'https://bifrost.example.com/openai/v1' ) assert _secret_value(result, 'llm.api_key') == 'sk-bf-shared-smoke-test' mock_client_class.assert_not_called() @pytest.mark.asyncio async def test_create_entries_local_deployment(self, mock_settings): """Test create_entries in local deployment mode.""" with patch.dict(os.environ, {'LOCAL_DEPLOYMENT': '1'}): with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch( 'storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com' ): result = await LiteLlmManager.create_entries( 'test-org-id', 'test-user-id', mock_settings, create_user=True ) assert result is not None assert _agent_value(result, 'agent') == 'CodeActAgent' assert ( _agent_value(result, 'llm.model') == get_default_litellm_model() ) assert _secret_value(result, 'llm.api_key') == 'test-key' assert _agent_value(result, 'llm.base_url') == 'http://test.com' @pytest.mark.asyncio async def test_create_entries_cloud_deployment(self, mock_settings, mock_response): """Test create_entries in cloud deployment mode.""" mock_404_response = MagicMock() mock_404_response.status_code = 404 mock_404_response.is_success = False mock_404_response.raise_for_status.side_effect = httpx.HTTPStatusError( message='Not Found', request=MagicMock(), response=mock_404_response ) # Mock user exists check response mock_user_exists_response = MagicMock() mock_user_exists_response.is_success = True mock_user_exists_response.json.return_value = { 'user_info': {'user_id': 'test-user-id'} } mock_token_manager = MagicMock() mock_token_manager.return_value.get_user_info_from_user_id = AsyncMock( return_value={'email': 'test@example.com'} ) mock_client = AsyncMock() # First GET is for _get_team (404), second GET is for _user_exists (success) mock_client.get.side_effect = [mock_404_response, mock_user_exists_response] mock_client.post.return_value = mock_response mock_client_class = MagicMock() mock_client_class.return_value.__aenter__.return_value = mock_client with ( patch.dict(os.environ, {'LOCAL_DEPLOYMENT': ''}), patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'), patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'), patch('storage.lite_llm_manager.TokenManager', mock_token_manager), patch('httpx.AsyncClient', mock_client_class), ): result = await LiteLlmManager.create_entries( 'test-org-id', 'test-user-id', mock_settings, create_user=False ) assert result is not None assert _agent_value(result, 'agent') == 'CodeActAgent' assert _agent_value(result, 'llm.model') == get_default_litellm_model() assert _secret_value(result, 'llm.api_key') == 'test-api-key' assert _agent_value(result, 'llm.base_url') == 'http://test.com' # Verify API calls were made (get_team + user_exists + 4 posts) assert mock_client.get.call_count == 2 # get_team + user_exists assert ( mock_client.post.call_count == 4 ) # create_team, add_user_to_team, delete_key_by_alias, generate_key @pytest.mark.asyncio async def test_create_entries_inherits_existing_team_budget( self, mock_settings, mock_response ): """Test that create_entries inherits budget from existing team.""" mock_team_response = MagicMock() mock_team_response.is_success = True mock_team_response.status_code = 200 mock_team_response.json.return_value = { 'team_info': {'max_budget': 30.0, 'spend': 5.0}, 'team_memberships': [], } mock_team_response.raise_for_status = MagicMock() # Mock user exists check response mock_user_exists_response = MagicMock() mock_user_exists_response.is_success = True mock_user_exists_response.json.return_value = { 'user_info': {'user_id': 'test-user-id'} } mock_token_manager = MagicMock() mock_token_manager.return_value.get_user_info_from_user_id = AsyncMock( return_value={'email': 'test@example.com'} ) mock_client = AsyncMock() # First GET is for _get_team (success), second GET is for _user_exists (success) mock_client.get.side_effect = [mock_team_response, mock_user_exists_response] mock_client.post.return_value = mock_response mock_client_class = MagicMock() mock_client_class.return_value.__aenter__.return_value = mock_client with ( patch.dict(os.environ, {'LOCAL_DEPLOYMENT': ''}), patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'), patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'), patch('storage.lite_llm_manager.TokenManager', mock_token_manager), patch('httpx.AsyncClient', mock_client_class), ): result = await LiteLlmManager.create_entries( 'test-org-id', 'test-user-id', mock_settings, create_user=False ) assert result is not None # Verify _get_team was called first assert mock_client.get.call_count == 2 # get_team + user_exists get_call_url = mock_client.get.call_args_list[0][0][0] assert 'team/info' in get_call_url assert 'test-org-id' in get_call_url # Verify _create_team was called with inherited budget (30.0) create_team_call = mock_client.post.call_args_list[0] assert 'team/new' in create_team_call[0][0] assert create_team_call[1]['json']['max_budget'] == 30.0 # Verify _add_user_to_team was called with inherited budget (30.0) add_user_call = mock_client.post.call_args_list[1] assert 'team/member_add' in add_user_call[0][0] assert add_user_call[1]['json']['max_budget_in_team'] == 30.0 @pytest.mark.asyncio async def test_create_entries_new_org_uses_default_initial_budget( self, mock_settings, mock_response ): """Test that create_entries uses DEFAULT_INITIAL_BUDGET for new org.""" mock_404_response = MagicMock() mock_404_response.status_code = 404 mock_404_response.is_success = False mock_404_response.raise_for_status.side_effect = httpx.HTTPStatusError( message='Not Found', request=MagicMock(), response=mock_404_response ) mock_token_manager = MagicMock() mock_token_manager.return_value.get_user_info_from_user_id = AsyncMock( return_value={'email': 'test@example.com'} ) # Mock user exists check response mock_user_exists_response = MagicMock() mock_user_exists_response.is_success = True mock_user_exists_response.json.return_value = { 'user_info': {'user_id': 'test-user-id'} } mock_client = AsyncMock() # First GET is for _get_team (404), second GET is for _user_exists (success) mock_client.get.side_effect = [mock_404_response, mock_user_exists_response] mock_client.post.return_value = mock_response mock_client_class = MagicMock() mock_client_class.return_value.__aenter__.return_value = mock_client with ( patch.dict(os.environ, {'LOCAL_DEPLOYMENT': ''}), patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'), patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'), patch('storage.lite_llm_manager.TokenManager', mock_token_manager), patch('httpx.AsyncClient', mock_client_class), patch('storage.lite_llm_manager.DEFAULT_INITIAL_BUDGET', 0.0), ): result = await LiteLlmManager.create_entries( 'test-org-id', 'test-user-id', mock_settings, create_user=False ) assert result is not None # Verify _create_team was called with DEFAULT_INITIAL_BUDGET (0.0) create_team_call = mock_client.post.call_args_list[0] assert 'team/new' in create_team_call[0][0] assert create_team_call[1]['json']['max_budget'] == 0.0 # Verify _add_user_to_team was called with DEFAULT_INITIAL_BUDGET (0.0) add_user_call = mock_client.post.call_args_list[1] assert 'team/member_add' in add_user_call[0][0] assert add_user_call[1]['json']['max_budget_in_team'] == 0.0 @pytest.mark.asyncio async def test_create_entries_new_org_uses_custom_default_budget( self, mock_settings, mock_response ): """Test that create_entries uses custom DEFAULT_INITIAL_BUDGET for new org.""" mock_404_response = MagicMock() mock_404_response.status_code = 404 mock_404_response.is_success = False mock_404_response.raise_for_status.side_effect = httpx.HTTPStatusError( message='Not Found', request=MagicMock(), response=mock_404_response ) # Mock user exists check response mock_user_exists_response = MagicMock() mock_user_exists_response.is_success = True mock_user_exists_response.json.return_value = { 'user_info': {'user_id': 'test-user-id'} } mock_token_manager = MagicMock() mock_token_manager.return_value.get_user_info_from_user_id = AsyncMock( return_value={'email': 'test@example.com'} ) mock_client = AsyncMock() # First GET is for _get_team (404), second GET is for _user_exists (success) mock_client.get.side_effect = [mock_404_response, mock_user_exists_response] mock_client.post.return_value = mock_response mock_client_class = MagicMock() mock_client_class.return_value.__aenter__.return_value = mock_client custom_budget = 50.0 with ( patch.dict(os.environ, {'LOCAL_DEPLOYMENT': ''}), patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'), patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'), patch('storage.lite_llm_manager.TokenManager', mock_token_manager), patch('httpx.AsyncClient', mock_client_class), patch('storage.lite_llm_manager.DEFAULT_INITIAL_BUDGET', custom_budget), ): result = await LiteLlmManager.create_entries( 'test-org-id', 'test-user-id', mock_settings, create_user=False ) assert result is not None # Verify _create_team was called with custom DEFAULT_INITIAL_BUDGET create_team_call = mock_client.post.call_args_list[0] assert 'team/new' in create_team_call[0][0] assert create_team_call[1]['json']['max_budget'] == custom_budget # Verify _add_user_to_team was called with custom DEFAULT_INITIAL_BUDGET add_user_call = mock_client.post.call_args_list[1] assert 'team/member_add' in add_user_call[0][0] assert add_user_call[1]['json']['max_budget_in_team'] == custom_budget @pytest.mark.asyncio async def test_create_entries_propagates_non_404_errors(self, mock_settings): """Test that create_entries propagates non-404 errors from _get_team.""" mock_500_response = MagicMock() mock_500_response.status_code = 500 mock_500_response.is_success = False mock_token_manager = MagicMock() mock_token_manager.return_value.get_user_info_from_user_id = AsyncMock( return_value={'email': 'test@example.com'} ) mock_client = AsyncMock() mock_client.get.return_value = mock_500_response mock_client.get.return_value.raise_for_status.side_effect = ( httpx.HTTPStatusError( message='Internal Server Error', request=MagicMock(), response=mock_500_response, ) ) mock_client_class = MagicMock() mock_client_class.return_value.__aenter__.return_value = mock_client with ( patch.dict(os.environ, {'LOCAL_DEPLOYMENT': ''}), patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'), patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'), patch('storage.lite_llm_manager.TokenManager', mock_token_manager), patch('httpx.AsyncClient', mock_client_class), ): with pytest.raises(httpx.HTTPStatusError) as exc_info: await LiteLlmManager.create_entries( 'test-org-id', 'test-user-id', mock_settings, create_user=False ) assert exc_info.value.response.status_code == 500 @pytest.mark.asyncio async def test_migrate_entries_missing_config(self, mock_user_settings): """Test migrate_entries when LiteLLM config is missing.""" with patch.dict(os.environ, {'LITE_LLM_API_KEY': '', 'LITE_LLM_API_URL': ''}): with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', None): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', None): result = await LiteLlmManager.migrate_entries( 'test-org-id', 'test-user-id', mock_user_settings, ) assert result is None @pytest.mark.asyncio async def test_migrate_entries_local_deployment(self, mock_user_settings): """Test migrate_entries in local deployment mode.""" with patch.dict(os.environ, {'LOCAL_DEPLOYMENT': '1'}): with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch( 'storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com' ): result = await LiteLlmManager.migrate_entries( 'test-org-id', 'test-user-id', mock_user_settings, ) # migrate_entries returns the user_settings unchanged assert result is not None effective_settings = result.to_settings() assert _agent_value(effective_settings, 'agent') == 'TestAgent' assert _agent_value(effective_settings, 'llm.model') == 'test-model' assert result.llm_api_key.get_secret_value() == 'test-key' assert ( _agent_value(effective_settings, 'llm.base_url') == 'http://test.com' ) @pytest.mark.asyncio async def test_migrate_entries_no_user_found(self, mock_user_settings): """Test migrate_entries when user is not found.""" with patch.dict(os.environ, {'LOCAL_DEPLOYMENT': ''}): with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch( 'storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com' ): with patch( 'storage.lite_llm_manager.TokenManager' ) as mock_token_manager: mock_token_manager.return_value.get_user_info_from_user_id = ( AsyncMock(return_value={'email': 'test@example.com'}) ) # Mock the _get_user method directly to return None with patch.object( LiteLlmManager, '_get_user', new_callable=AsyncMock ) as mock_get_user: mock_get_user.return_value = None result = await LiteLlmManager.migrate_entries( 'test-org-id', 'test-user-id', mock_user_settings, ) assert result is None @pytest.mark.asyncio async def test_migrate_entries_already_migrated( self, mock_user_settings, mock_user_response ): """Test migrate_entries when user is already migrated (no max_budget).""" mock_user_response.json.return_value = { 'user_info': { 'max_budget': None, # Already migrated 'spend': 10.0, } } with patch.dict(os.environ, {'LOCAL_DEPLOYMENT': ''}): with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch( 'storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com' ): with patch( 'storage.lite_llm_manager.TokenManager' ) as mock_token_manager: mock_token_manager.return_value.get_user_info_from_user_id = ( AsyncMock(return_value={'email': 'test@example.com'}) ) with patch('httpx.AsyncClient') as mock_client_class: mock_client = AsyncMock() mock_client_class.return_value.__aenter__.return_value = ( mock_client ) mock_client.get.return_value = mock_user_response result = await LiteLlmManager.migrate_entries( 'test-org-id', 'test-user-id', mock_user_settings, ) assert result is None @pytest.mark.asyncio async def test_migrate_entries_successful_migration( self, mock_user_settings, mock_user_response, mock_response ): """Test successful migrate_entries operation.""" # Mock response for key list mock_key_list_response = MagicMock() mock_key_list_response.is_success = True mock_key_list_response.status_code = 200 mock_key_list_response.json.return_value = { 'keys': ['test-key-1', 'test-key-2'], 'total_count': 2, } mock_key_list_response.raise_for_status = MagicMock() with patch.dict(os.environ, {'LOCAL_DEPLOYMENT': ''}): with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch( 'storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com' ): with patch( 'storage.lite_llm_manager.TokenManager' ) as mock_token_manager: mock_token_manager.return_value.get_user_info_from_user_id = ( AsyncMock(return_value={'email': 'test@example.com'}) ) with patch('httpx.AsyncClient') as mock_client_class: mock_client = AsyncMock() mock_client_class.return_value.__aenter__.return_value = ( mock_client ) # First GET is for _get_user, second GET is for _get_user_keys mock_client.get.side_effect = [ mock_user_response, mock_key_list_response, ] mock_client.post.return_value = mock_response # Mock verify_key to return True (key exists in LiteLLM) with patch.object( LiteLlmManager, 'verify_key', return_value=True ): result = await LiteLlmManager.migrate_entries( 'test-org-id', 'test-user-id', mock_user_settings, ) # migrate_entries returns the user_settings unchanged assert result is not None effective_settings = result.to_settings() assert ( _agent_value(effective_settings, 'agent') == 'TestAgent' ) assert ( _agent_value(effective_settings, 'llm.model') == 'test-model' ) assert result.llm_api_key.get_secret_value() == 'test-key' assert ( _agent_value(effective_settings, 'llm.base_url') == 'http://test.com' ) # Verify migration steps were called: # - 2 GET requests: _get_user, _get_user_keys # - POST requests: create_team, update_user, add_user_to_team, # and update_key for each key (2 keys) assert mock_client.get.call_count == 2 assert ( mock_client.post.call_count == 5 ) # create_team, update_user, add_user_to_team, 2x update_key @pytest.mark.asyncio async def test_migrate_entries_generates_key_when_db_key_not_in_litellm( self, mock_user_settings, mock_user_response, mock_response ): """Test migrate_entries generates a new key when the DB key doesn't exist in LiteLLM.""" # Mock response for key list mock_key_list_response = MagicMock() mock_key_list_response.is_success = True mock_key_list_response.status_code = 200 mock_key_list_response.json.return_value = { 'keys': ['test-key-1', 'test-key-2'], 'total_count': 2, } mock_key_list_response.raise_for_status = MagicMock() # Mock response for key generation mock_generate_response = MagicMock() mock_generate_response.is_success = True mock_generate_response.status_code = 200 mock_generate_response.json.return_value = {'key': 'new-generated-key'} mock_generate_response.raise_for_status = MagicMock() with patch.dict(os.environ, {'LOCAL_DEPLOYMENT': ''}): with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch( 'storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com' ): with patch( 'storage.lite_llm_manager.TokenManager' ) as mock_token_manager: mock_token_manager.return_value.get_user_info_from_user_id = ( AsyncMock(return_value={'email': 'test@example.com'}) ) with patch('httpx.AsyncClient') as mock_client_class: mock_client = AsyncMock() mock_client_class.return_value.__aenter__.return_value = ( mock_client ) # First GET is for _get_user, second GET is for _get_user_keys mock_client.get.side_effect = [ mock_user_response, mock_key_list_response, ] # POST responses: create_team, update_user, add_user_to_team, # 2x update_key, and 1x generate_key mock_client.post.side_effect = [ mock_response, # create_team mock_response, # update_user mock_response, # add_user_to_team mock_response, # update_key 1 mock_response, # update_key 2 mock_generate_response, # generate_key ] # Mock verify_key to return False (key doesn't exist in LiteLLM) with patch.object( LiteLlmManager, 'verify_key', return_value=False ): result = await LiteLlmManager.migrate_entries( 'test-org-id', 'test-user-id', mock_user_settings, ) # migrate_entries should update user_settings with the new key assert result is not None assert ( result.agent_settings['llm']['api_key'] == 'new-generated-key' ) assert result.llm_api_key_for_byor_secret is not None assert ( result.llm_api_key_for_byor_secret.get_secret_value() == 'new-generated-key' ) # Verify migration steps were called including key generation: # - 2 GET requests: _get_user, _get_user_keys # - 6 POST requests: create_team, update_user, add_user_to_team, # 2x update_key, 1x generate_key assert mock_client.get.call_count == 2 assert mock_client.post.call_count == 6 @pytest.mark.asyncio async def test_update_team_and_users_budget_missing_config(self): """Test update_team_and_users_budget when LiteLLM config is missing.""" with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', None): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', None): # Should not raise an exception, just return early await LiteLlmManager.update_team_and_users_budget('test-team-id', 100.0) @pytest.mark.asyncio async def test_update_team_and_users_budget_successful( self, mock_team_response, mock_response ): """Test successful update_team_and_users_budget operation.""" with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): with patch('httpx.AsyncClient') as mock_client_class: mock_client = AsyncMock() mock_client_class.return_value.__aenter__.return_value = mock_client mock_client.post.return_value = mock_response mock_client.get.return_value = mock_team_response await LiteLlmManager.update_team_and_users_budget( 'test-team-id', 100.0 ) # Verify update_team and update_user_in_team were called assert ( mock_client.post.call_count == 2 ) # update_team, update_user_in_team @pytest.mark.asyncio async def test_create_team_success(self, mock_http_client, mock_response): """Test successful _create_team operation.""" mock_http_client.post.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): await LiteLlmManager._create_team( mock_http_client, 'test-alias', 'test-team-id', 100.0 ) mock_http_client.post.assert_called_once() call_args = mock_http_client.post.call_args assert 'http://test.com/team/new' in call_args[0] assert call_args[1]['json']['team_id'] == 'test-team-id' assert call_args[1]['json']['team_alias'] == 'test-alias' assert call_args[1]['json']['max_budget'] == 100.0 @pytest.mark.asyncio async def test_create_team_already_exists(self, mock_http_client): """Test _create_team when team already exists.""" error_response = MagicMock() error_response.is_success = False error_response.status_code = 400 error_response.text = 'Team already exists. Please use a different team id' mock_http_client.post.return_value = error_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): with patch.object( LiteLlmManager, '_update_team', new_callable=AsyncMock ) as mock_update: await LiteLlmManager._create_team( mock_http_client, 'test-alias', 'test-team-id', 100.0 ) mock_update.assert_called_once_with( mock_http_client, 'test-team-id', 'test-alias', 100.0 ) @pytest.mark.asyncio async def test_create_team_error(self, mock_http_client): """Test _create_team with unexpected error.""" error_response = MagicMock() error_response.is_success = False error_response.status_code = 500 error_response.text = 'Internal server error' error_response.raise_for_status.side_effect = httpx.HTTPStatusError( 'Server error', request=MagicMock(), response=error_response ) mock_http_client.post.return_value = error_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): with pytest.raises(httpx.HTTPStatusError): await LiteLlmManager._create_team( mock_http_client, 'test-alias', 'test-team-id', 100.0 ) @pytest.mark.asyncio async def test_get_team_success(self, mock_http_client, mock_team_response): """Test successful _get_team operation.""" mock_http_client.get.return_value = mock_team_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): result = await LiteLlmManager._get_team( mock_http_client, 'test-team-id' ) assert result is not None assert 'team_memberships' in result mock_http_client.get.assert_called_once_with( 'http://test.com/team/info?team_id=test-team-id' ) @pytest.mark.asyncio async def test_create_user_success(self, mock_http_client, mock_response): """Test successful _create_user operation returns True.""" mock_http_client.post.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): result = await LiteLlmManager._create_user( mock_http_client, 'test@example.com', 'test-user-id' ) assert result is True mock_http_client.post.assert_called_once() call_args = mock_http_client.post.call_args assert 'http://test.com/user/new' in call_args[0] assert call_args[1]['json']['user_email'] == 'test@example.com' assert call_args[1]['json']['user_id'] == 'test-user-id' @pytest.mark.asyncio async def test_create_user_duplicate_email(self, mock_http_client, mock_response): """Test _create_user with duplicate email handling returns True.""" # First call fails with duplicate email error_response = MagicMock() error_response.is_success = False error_response.status_code = 400 error_response.text = 'duplicate email' # Second call succeeds mock_http_client.post.side_effect = [error_response, mock_response] with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): result = await LiteLlmManager._create_user( mock_http_client, 'test@example.com', 'test-user-id' ) assert result is True assert mock_http_client.post.call_count == 2 # Second call should have None email second_call_args = mock_http_client.post.call_args_list[1] assert second_call_args[1]['json']['user_email'] is None @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_user_exists_returns_true(self, mock_http_client): """Test _user_exists returns True when user exists in LiteLLM.""" # Arrange user_response = MagicMock() user_response.is_success = True user_response.json.return_value = { 'user_info': {'user_id': 'test-user-id', 'email': 'test@example.com'} } mock_http_client.get.return_value = user_response # Act result = await LiteLlmManager._user_exists(mock_http_client, 'test-user-id') # Assert assert result is True mock_http_client.get.assert_called_once() @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_user_exists_returns_false_when_not_found(self, mock_http_client): """Test _user_exists returns False when user not found.""" # Arrange user_response = MagicMock() user_response.is_success = False mock_http_client.get.return_value = user_response # Act result = await LiteLlmManager._user_exists(mock_http_client, 'test-user-id') # Assert assert result is False @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_user_exists_returns_false_on_mismatched_user_id( self, mock_http_client ): """Test _user_exists returns False when returned user_id doesn't match.""" # Arrange user_response = MagicMock() user_response.is_success = True user_response.json.return_value = { 'user_info': {'user_id': 'different-user-id'} } mock_http_client.get.return_value = user_response # Act result = await LiteLlmManager._user_exists(mock_http_client, 'test-user-id') # Assert assert result is False @pytest.mark.asyncio @patch('storage.lite_llm_manager.logger') @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_create_user_already_exists_and_verified( self, mock_logger, mock_http_client ): """Test _create_user returns True when user already exists and is verified.""" # Arrange first_response = MagicMock() first_response.is_success = False first_response.status_code = 400 first_response.text = 'duplicate email' second_response = MagicMock() second_response.is_success = False second_response.status_code = 409 second_response.text = 'User with id test-user-id already exists' user_exists_response = MagicMock() user_exists_response.is_success = True user_exists_response.json.return_value = { 'user_info': {'user_id': 'test-user-id'} } mock_http_client.post.side_effect = [first_response, second_response] mock_http_client.get.return_value = user_exists_response # Act result = await LiteLlmManager._create_user( mock_http_client, 'test@example.com', 'test-user-id' ) # Assert assert result is True mock_logger.warning.assert_any_call( 'litellm_user_already_exists', extra={'user_id': 'test-user-id'}, ) @pytest.mark.asyncio @patch('storage.lite_llm_manager.logger') @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_create_user_already_exists_but_not_found_returns_false( self, mock_logger, mock_http_client ): """Test _create_user returns False when LiteLLM claims user exists but verification fails.""" # Arrange first_response = MagicMock() first_response.is_success = False first_response.status_code = 400 first_response.text = 'duplicate email' second_response = MagicMock() second_response.is_success = False second_response.status_code = 409 second_response.text = 'User with id test-user-id already exists' user_not_exists_response = MagicMock() user_not_exists_response.is_success = False mock_http_client.post.side_effect = [first_response, second_response] mock_http_client.get.return_value = user_not_exists_response # Act result = await LiteLlmManager._create_user( mock_http_client, 'test@example.com', 'test-user-id' ) # Assert assert result is False mock_logger.error.assert_any_call( 'litellm_user_claimed_exists_but_not_found', extra={ 'user_id': 'test-user-id', 'status_code': 409, 'text': 'User with id test-user-id already exists', }, ) @pytest.mark.asyncio @patch('storage.lite_llm_manager.logger') @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_create_user_failure_returns_false( self, mock_logger, mock_http_client ): """Test _create_user returns False when creation fails with non-'already exists' error.""" # Arrange first_response = MagicMock() first_response.is_success = False first_response.status_code = 400 first_response.text = 'duplicate email' second_response = MagicMock() second_response.is_success = False second_response.status_code = 500 second_response.text = 'Internal server error' mock_http_client.post.side_effect = [first_response, second_response] # Act result = await LiteLlmManager._create_user( mock_http_client, 'test@example.com', 'test-user-id' ) # Assert assert result is False mock_logger.error.assert_any_call( 'error_creating_litellm_user', extra={ 'status_code': 500, 'text': 'Internal server error', 'user_id': 'test-user-id', 'email': None, }, ) @pytest.mark.asyncio @patch('storage.lite_llm_manager.logger') @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_create_user_already_exists_with_409_status_code( self, mock_logger, mock_http_client ): """Test _create_user handles 409 Conflict when user already exists and verifies.""" # Arrange first_response = MagicMock() first_response.is_success = False first_response.status_code = 400 first_response.text = 'duplicate email' second_response = MagicMock() second_response.is_success = False second_response.status_code = 409 second_response.text = 'User with id test-user-id already exists' user_exists_response = MagicMock() user_exists_response.is_success = True user_exists_response.json.return_value = { 'user_info': {'user_id': 'test-user-id'} } mock_http_client.post.side_effect = [first_response, second_response] mock_http_client.get.return_value = user_exists_response # Act result = await LiteLlmManager._create_user( mock_http_client, 'test@example.com', 'test-user-id' ) # Assert assert result is True mock_logger.warning.assert_any_call( 'litellm_user_already_exists', extra={'user_id': 'test-user-id'}, ) @pytest.mark.asyncio @patch('storage.lite_llm_manager.logger') @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_create_user_already_exists_with_400_status_code( self, mock_logger, mock_http_client ): """Test _create_user handles 400 Bad Request when user already exists and verifies.""" # Arrange first_response = MagicMock() first_response.is_success = False first_response.status_code = 400 first_response.text = 'duplicate email' second_response = MagicMock() second_response.is_success = False second_response.status_code = 400 second_response.text = 'User already exists' user_exists_response = MagicMock() user_exists_response.is_success = True user_exists_response.json.return_value = { 'user_info': {'user_id': 'test-user-id'} } mock_http_client.post.side_effect = [first_response, second_response] mock_http_client.get.return_value = user_exists_response # Act result = await LiteLlmManager._create_user( mock_http_client, 'test@example.com', 'test-user-id' ) # Assert assert result is True mock_logger.warning.assert_any_call( 'litellm_user_already_exists', extra={'user_id': 'test-user-id'}, ) @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_add_user_to_team_success(self, mock_http_client, mock_response): """Test successful _add_user_to_team operation.""" # Arrange mock_http_client.post.return_value = mock_response # Act await LiteLlmManager._add_user_to_team( mock_http_client, 'test-user-id', 'test-team-id', 100.0 ) # Assert mock_http_client.post.assert_called_once() call_args = mock_http_client.post.call_args assert 'http://test.com/team/member_add' in call_args[0] assert call_args[1]['json']['team_id'] == 'test-team-id' assert call_args[1]['json']['member'] == { 'user_id': 'test-user-id', 'role': 'user', } assert call_args[1]['json']['max_budget_in_team'] == 100.0 @pytest.mark.asyncio @patch('storage.lite_llm_manager.logger') @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_add_user_to_team_already_in_team( self, mock_logger, mock_http_client ): """Test _add_user_to_team handles 'already in team' error gracefully.""" # Arrange error_response = MagicMock() error_response.is_success = False error_response.status_code = 400 error_response.text = ( '{"error":{"message":"User already in team. Member: ' 'user_id=test-user-id","type":"team_member_already_in_team"}}' ) mock_http_client.post.return_value = error_response # Act await LiteLlmManager._add_user_to_team( mock_http_client, 'test-user-id', 'test-team-id', 100.0 ) # Assert mock_logger.warning.assert_called_once_with( 'user_already_in_team', extra={ 'user_id': 'test-user-id', 'team_id': 'test-team-id', }, ) @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_add_user_to_team_other_error_raises_exception( self, mock_http_client ): """Test _add_user_to_team raises exception for non-'already in team' errors.""" # Arrange error_response = MagicMock() error_response.is_success = False error_response.status_code = 500 error_response.text = 'Internal server error' error_response.raise_for_status.side_effect = httpx.HTTPStatusError( 'Server error', request=MagicMock(), response=error_response ) mock_http_client.post.return_value = error_response # Act & Assert with pytest.raises(httpx.HTTPStatusError): await LiteLlmManager._add_user_to_team( mock_http_client, 'test-user-id', 'test-team-id', 100.0 ) @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_update_key_success(self, mock_http_client, mock_response): """Test successful _update_key operation.""" # Arrange mock_http_client.post.return_value = mock_response # Act await LiteLlmManager._update_key( mock_http_client, 'test-user-id', 'test-api-key', team_id='test-team-id' ) # Assert mock_http_client.post.assert_called_once() call_args = mock_http_client.post.call_args assert 'http://test.com/key/update' in call_args[0] assert call_args[1]['json']['key'] == 'test-api-key' assert call_args[1]['json']['team_id'] == 'test-team-id' @pytest.mark.asyncio @patch('storage.lite_llm_manager.logger') @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_update_key_invalid_key_returns_gracefully( self, mock_logger, mock_http_client ): """Test _update_key handles 401 Unauthorized for invalid keys gracefully.""" # Arrange error_response = MagicMock() error_response.is_success = False error_response.status_code = 401 error_response.text = 'Unauthorized' mock_http_client.post.return_value = error_response # Act await LiteLlmManager._update_key( mock_http_client, 'test-user-id', 'invalid-api-key', team_id='test-team-id' ) # Assert mock_logger.warning.assert_called_once_with( 'invalid_litellm_key_during_update', extra={'user_id': 'test-user-id', 'text': 'Unauthorized'}, ) @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_update_key_other_error_raises_exception(self, mock_http_client): """Test _update_key raises exception for non-401 errors.""" # Arrange error_response = MagicMock() error_response.is_success = False error_response.status_code = 500 error_response.text = 'Internal server error' error_response.raise_for_status.side_effect = httpx.HTTPStatusError( 'Server error', request=MagicMock(), response=error_response ) mock_http_client.post.return_value = error_response # Act & Assert with pytest.raises(httpx.HTTPStatusError): await LiteLlmManager._update_key( mock_http_client, 'test-user-id', 'test-api-key', team_id='test-team-id' ) @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_get_user_keys_success(self, mock_http_client): """Test successful _get_user_keys operation.""" # Arrange mock_response = MagicMock() mock_response.is_success = True mock_response.status_code = 200 mock_response.json.return_value = { 'keys': ['key-1', 'key-2', 'key-3'], 'total_count': 3, } mock_http_client.get.return_value = mock_response # Act keys = await LiteLlmManager._get_user_keys(mock_http_client, 'test-user-id') # Assert assert keys == ['key-1', 'key-2', 'key-3'] mock_http_client.get.assert_called_once() call_args = mock_http_client.get.call_args assert 'http://test.com/key/list' in call_args[0] assert call_args[1]['params'] == {'user_id': 'test-user-id'} @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_get_user_keys_empty_list(self, mock_http_client): """Test _get_user_keys returns empty list when user has no keys.""" # Arrange mock_response = MagicMock() mock_response.is_success = True mock_response.status_code = 200 mock_response.json.return_value = { 'keys': [], 'total_count': 0, } mock_http_client.get.return_value = mock_response # Act keys = await LiteLlmManager._get_user_keys(mock_http_client, 'test-user-id') # Assert assert keys == [] @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_get_user_keys_error_returns_empty_list(self, mock_http_client): """Test _get_user_keys returns empty list on error.""" # Arrange mock_response = MagicMock() mock_response.is_success = False mock_response.status_code = 500 mock_response.text = 'Internal server error' mock_http_client.get.return_value = mock_response # Act keys = await LiteLlmManager._get_user_keys(mock_http_client, 'test-user-id') # Assert assert keys == [] @pytest.mark.asyncio async def test_get_user_keys_missing_config(self, mock_http_client): """Test _get_user_keys returns empty list when config is missing.""" with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', None): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', None): keys = await LiteLlmManager._get_user_keys( mock_http_client, 'test-user-id' ) assert keys == [] @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_update_user_keys_success(self, mock_http_client, mock_response): """Test successful _update_user_keys operation.""" # Arrange mock_key_list_response = MagicMock() mock_key_list_response.is_success = True mock_key_list_response.status_code = 200 mock_key_list_response.json.return_value = { 'keys': ['key-1', 'key-2'], 'total_count': 2, } mock_http_client.get.return_value = mock_key_list_response mock_http_client.post.return_value = mock_response # Act await LiteLlmManager._update_user_keys( mock_http_client, 'test-user-id', team_id='test-team-id' ) # Assert # Should call GET once for key list assert mock_http_client.get.call_count == 1 # Should call POST twice (once for each key) assert mock_http_client.post.call_count == 2 @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_update_user_keys_no_keys(self, mock_http_client, mock_response): """Test _update_user_keys when user has no keys.""" # Arrange mock_key_list_response = MagicMock() mock_key_list_response.is_success = True mock_key_list_response.status_code = 200 mock_key_list_response.json.return_value = { 'keys': [], 'total_count': 0, } mock_http_client.get.return_value = mock_key_list_response # Act await LiteLlmManager._update_user_keys( mock_http_client, 'test-user-id', team_id='test-team-id' ) # Assert # Should call GET once for key list assert mock_http_client.get.call_count == 1 # Should not call POST since there are no keys assert mock_http_client.post.call_count == 0 @pytest.mark.asyncio async def test_generate_key_success(self, mock_http_client, mock_response): """Test successful _generate_key operation.""" mock_http_client.post.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): result = await LiteLlmManager._generate_key( mock_http_client, 'test-user-id', 'test-team-id', 'test-alias', {'test': 'metadata'}, ) assert result == 'test-api-key' mock_http_client.post.assert_called_once() call_args = mock_http_client.post.call_args assert 'http://test.com/key/generate' in call_args[0] assert call_args[1]['json']['user_id'] == 'test-user-id' assert call_args[1]['json']['team_id'] == 'test-team-id' assert call_args[1]['json']['key_alias'] == 'test-alias' assert call_args[1]['json']['metadata'] == {'test': 'metadata'} @pytest.mark.asyncio async def test_get_key_info_success(self, mock_http_client, mock_key_info_response): """Test successful _get_key_info operation.""" mock_http_client.get.return_value = mock_key_info_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): with patch('storage.user_store.UserStore') as mock_user_store: # Mock user with org member mock_user = MagicMock() mock_org_member = MagicMock() mock_org_member.org_id = 'test-ord-id' mock_org_member.llm_api_key = 'test-api-key' mock_user.org_members = [mock_org_member] mock_user_store.get_user_by_id = AsyncMock(return_value=mock_user) result = await LiteLlmManager._get_key_info( mock_http_client, 'test-ord-id', 'test-user-id' ) assert result is not None assert result['key_max_budget'] == 100.0 assert result['key_spend'] == 25.0 @pytest.mark.asyncio async def test_get_key_info_no_user(self, mock_http_client): """Test _get_key_info when user is not found.""" with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): with patch('storage.user_store.UserStore') as mock_user_store: mock_user_store.get_user_by_id = AsyncMock(return_value=None) result = await LiteLlmManager._get_key_info( mock_http_client, 'test-ord-id', 'test-user-id' ) assert result == {} @pytest.mark.asyncio async def test_delete_key_success(self, mock_http_client, mock_response): """Test successful _delete_key operation.""" mock_http_client.post.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): await LiteLlmManager._delete_key(mock_http_client, 'test-key-id') mock_http_client.post.assert_called_once() call_args = mock_http_client.post.call_args assert 'http://test.com/key/delete' in call_args[0] assert call_args[1]['json']['keys'] == ['test-key-id'] @pytest.mark.asyncio async def test_delete_key_not_found(self, mock_http_client): """Test _delete_key when key is not found (404 error).""" error_response = MagicMock() error_response.is_success = False error_response.status_code = 404 error_response.text = 'Key not found' mock_http_client.post.return_value = error_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): # Should not raise an exception for 404 await LiteLlmManager._delete_key(mock_http_client, 'test-key-id') @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_delete_key_not_found_with_alias_triggers_alias_deletion( self, mock_http_client ): """Test _delete_key falls back to alias deletion when key_id returns 404.""" # Arrange not_found_response = MagicMock() not_found_response.is_success = False not_found_response.status_code = 404 not_found_response.text = 'Key not found' alias_success_response = MagicMock() alias_success_response.is_success = True alias_success_response.status_code = 200 mock_http_client.post.side_effect = [not_found_response, alias_success_response] # Act await LiteLlmManager._delete_key( mock_http_client, 'test-key-id', key_alias='BYOR Key - user 123, org 456' ) # Assert assert mock_http_client.post.call_count == 2 first_call = mock_http_client.post.call_args_list[0] assert first_call[1]['json']['keys'] == ['test-key-id'] second_call = mock_http_client.post.call_args_list[1] assert second_call[1]['json']['key_aliases'] == ['BYOR Key - user 123, org 456'] @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_delete_key_not_found_without_alias_no_fallback( self, mock_http_client ): """Test _delete_key without alias does not attempt alias deletion on 404.""" # Arrange not_found_response = MagicMock() not_found_response.is_success = False not_found_response.status_code = 404 not_found_response.text = 'Key not found' mock_http_client.post.return_value = not_found_response # Act await LiteLlmManager._delete_key(mock_http_client, 'test-key-id') # Assert assert mock_http_client.post.call_count == 1 @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_delete_key_by_alias_success(self, mock_http_client, mock_response): """Test successful _delete_key_by_alias operation.""" # Arrange mock_http_client.post.return_value = mock_response # Act await LiteLlmManager._delete_key_by_alias( mock_http_client, 'BYOR Key - user 123, org 456' ) # Assert mock_http_client.post.assert_called_once() call_args = mock_http_client.post.call_args assert 'http://test.com/key/delete' in call_args[0] assert call_args[1]['json']['key_aliases'] == ['BYOR Key - user 123, org 456'] @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_delete_key_by_alias_not_found(self, mock_http_client): """Test _delete_key_by_alias when alias is not found (404).""" # Arrange not_found_response = MagicMock() not_found_response.is_success = False not_found_response.status_code = 404 not_found_response.text = 'Key alias not found' mock_http_client.post.return_value = not_found_response # Act & Assert - should not raise exception for 404 await LiteLlmManager._delete_key_by_alias( mock_http_client, 'BYOR Key - user 123, org 456' ) @pytest.mark.asyncio @patch('storage.lite_llm_manager.logger') @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') async def test_delete_key_by_alias_server_error_logs_warning( self, mock_logger, mock_http_client ): """Test _delete_key_by_alias logs warning for non-404 errors.""" # Arrange error_response = MagicMock() error_response.is_success = False error_response.status_code = 500 error_response.text = 'Internal server error' mock_http_client.post.return_value = error_response # Act await LiteLlmManager._delete_key_by_alias( mock_http_client, 'BYOR Key - user 123, org 456' ) # Assert mock_logger.warning.assert_called_once() call_args = mock_logger.warning.call_args assert call_args[0][0] == 'error_deleting_key_by_alias' @pytest.mark.asyncio @patch('storage.lite_llm_manager.LITE_LLM_API_URL', None) @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', None) async def test_delete_key_by_alias_missing_config(self, mock_http_client): """Test _delete_key_by_alias returns early when config is missing.""" # Act await LiteLlmManager._delete_key_by_alias( mock_http_client, 'BYOR Key - user 123, org 456' ) # Assert mock_http_client.post.assert_not_called() @pytest.mark.asyncio async def test_with_http_client_decorator(self): """Test the with_http_client decorator functionality.""" # Create a mock internal function async def mock_internal_fn(client, arg1, arg2, kwarg1=None): return f'client={type(client).__name__}, arg1={arg1}, arg2={arg2}, kwarg1={kwarg1}' # Apply the decorator decorated_fn = LiteLlmManager.with_http_client(mock_internal_fn) with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('httpx.AsyncClient') as mock_client_class: mock_client = AsyncMock() mock_client_class.return_value.__aenter__.return_value = mock_client result = await decorated_fn('test1', 'test2', kwarg1='test3') # Verify the client was injected as the first argument assert 'client=AsyncMock' in result assert 'arg1=test1' in result assert 'arg2=test2' in result assert 'kwarg1=test3' in result def test_public_methods_exist(self): """Test that all public wrapper methods exist and are properly decorated.""" public_methods = [ 'create_team', 'get_team', 'update_team', 'create_user', 'get_user', 'update_user', 'delete_user', 'add_user_to_team', 'get_user_team_info', 'update_user_in_team', 'generate_key', 'get_key_info', 'delete_key', ] for method_name in public_methods: assert hasattr(LiteLlmManager, method_name) method = getattr(LiteLlmManager, method_name) assert callable(method) # The methods are created by the with_http_client decorator, so they're functions # We can verify they exist and are callable, which is the important part @pytest.mark.asyncio async def test_error_handling_missing_config_all_methods(self): """Test that all methods handle missing configuration gracefully.""" with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', None): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', None): mock_client = AsyncMock() # Test all private methods that check for config await LiteLlmManager._create_team( mock_client, 'alias', 'team_id', 100.0 ) await LiteLlmManager._update_team( mock_client, 'team_id', 'alias', 100.0 ) await LiteLlmManager._create_user(mock_client, 'email', 'user_id') await LiteLlmManager._update_user(mock_client, 'user_id') await LiteLlmManager._delete_user(mock_client, 'user_id') await LiteLlmManager._add_user_to_team( mock_client, 'user_id', 'team_id', 100.0 ) await LiteLlmManager._update_user_in_team( mock_client, 'user_id', 'team_id', 100.0 ) await LiteLlmManager._delete_key(mock_client, 'key_id') result1 = await LiteLlmManager._get_team(mock_client, 'team_id') result2 = await LiteLlmManager._get_user(mock_client, 'user_id') # _generate_key raises ValueError when config is missing with pytest.raises( ValueError, match='LiteLLM API configuration not found' ): await LiteLlmManager._generate_key( mock_client, 'user_id', 'team_id', 'alias', {} ) result4 = await LiteLlmManager._get_user_team_info( mock_client, 'user_id', 'team_id' ) result5 = await LiteLlmManager._get_key_info( mock_client, 'test-ord-id', 'user_id' ) # Methods that return None when config is missing assert result1 is None assert result2 is None assert result4 is None assert result5 is None # Verify no HTTP calls were made mock_client.get.assert_not_called() mock_client.post.assert_not_called() @pytest.mark.asyncio async def test_delete_team_success(self, mock_http_client, mock_response): """ GIVEN: Valid team_id and configured LiteLLM API WHEN: delete_team is called THEN: Team is deleted successfully via POST /team/delete """ # Arrange team_id = 'test-team-123' mock_response.is_success = True mock_response.status_code = 200 mock_http_client.post.return_value = mock_response with ( patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'), patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.url'), patch('storage.lite_llm_manager.LITE_LLM_TEAM_ID', 'test-team'), ): # Act await LiteLlmManager._delete_team(mock_http_client, team_id) # Assert mock_http_client.post.assert_called_once_with( 'http://test.url/team/delete', json={'team_ids': [team_id]}, ) @pytest.mark.asyncio async def test_delete_team_not_found_is_idempotent( self, mock_http_client, mock_response ): """ GIVEN: Team does not exist (404 response) WHEN: delete_team is called THEN: Operation succeeds without raising exception (idempotent) """ # Arrange team_id = 'non-existent-team' mock_response.is_success = False mock_response.status_code = 404 mock_http_client.post.return_value = mock_response with ( patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'), patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.url'), patch('storage.lite_llm_manager.LITE_LLM_TEAM_ID', 'test-team'), ): # Act - should not raise await LiteLlmManager._delete_team(mock_http_client, team_id) # Assert mock_http_client.post.assert_called_once() @pytest.mark.asyncio async def test_delete_team_api_error_raises_exception( self, mock_http_client, mock_response ): """ GIVEN: LiteLLM API returns error (non-404) WHEN: delete_team is called THEN: HTTPStatusError is raised """ # Arrange team_id = 'test-team-123' mock_response.is_success = False mock_response.status_code = 500 mock_response.text = 'Internal Server Error' mock_response.raise_for_status = MagicMock( side_effect=httpx.HTTPStatusError( 'Server error', request=MagicMock(), response=mock_response ) ) mock_http_client.post.return_value = mock_response with ( patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'), patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.url'), patch('storage.lite_llm_manager.LITE_LLM_TEAM_ID', 'test-team'), ): # Act & Assert with pytest.raises(httpx.HTTPStatusError): await LiteLlmManager._delete_team(mock_http_client, team_id) @pytest.mark.asyncio async def test_delete_team_no_config_returns_early(self, mock_http_client): """ GIVEN: LiteLLM API is not configured WHEN: delete_team is called THEN: Function returns early without making API call """ # Arrange team_id = 'test-team-123' with ( patch('storage.lite_llm_manager.LITE_LLM_API_KEY', None), patch('storage.lite_llm_manager.LITE_LLM_API_URL', None), ): # Act await LiteLlmManager._delete_team(mock_http_client, team_id) # Assert mock_http_client.post.assert_not_called() @pytest.mark.asyncio async def test_delete_team_public_method(self): """ GIVEN: Valid team_id WHEN: Public delete_team method is called THEN: HTTP client is created and team is deleted """ # Arrange team_id = 'test-team-123' mock_response = AsyncMock() mock_response.is_success = True mock_response.status_code = 200 with ( patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'), patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.url'), patch('storage.lite_llm_manager.LITE_LLM_TEAM_ID', 'test-team'), patch('httpx.AsyncClient') as mock_client_class, ): mock_client = AsyncMock() mock_client.post.return_value = mock_response mock_client_class.return_value.__aenter__.return_value = mock_client # Act await LiteLlmManager.delete_team(team_id) # Assert mock_client.post.assert_called_once_with( 'http://test.url/team/delete', json={'team_ids': [team_id]}, ) @pytest.mark.asyncio async def test_remove_user_from_team_successful(self): """ GIVEN: Valid user_id and team_id WHEN: _remove_user_from_team is called THEN: HTTP POST is made to remove user from team """ mock_response = AsyncMock() mock_response.is_success = True mock_response.status_code = 200 with ( patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'), patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.url'), ): mock_client = AsyncMock() mock_client.post.return_value = mock_response await LiteLlmManager._remove_user_from_team( mock_client, 'test-user-id', 'test-team-id' ) mock_client.post.assert_called_once_with( 'http://test.url/team/member_delete', json={ 'team_id': 'test-team-id', 'user_id': 'test-user-id', }, ) @pytest.mark.asyncio async def test_remove_user_from_team_not_found(self): """ GIVEN: User not in team WHEN: _remove_user_from_team is called THEN: 404 response is handled gracefully without raising """ mock_response = AsyncMock() mock_response.is_success = False mock_response.status_code = 404 mock_response.text = 'User not found in team' mock_response.raise_for_status = MagicMock() with ( patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'), patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.url'), ): mock_client = AsyncMock() mock_client.post.return_value = mock_response # Should not raise an exception await LiteLlmManager._remove_user_from_team( mock_client, 'test-user-id', 'test-team-id' ) @pytest.mark.asyncio async def test_downgrade_entries_missing_config(self, mock_user_settings): """Test downgrade_entries when LiteLLM config is missing.""" with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', None): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', None): result = await LiteLlmManager.downgrade_entries( 'test-org-id', 'test-user-id', mock_user_settings, ) assert result is None @pytest.mark.asyncio async def test_downgrade_entries_team_not_found(self, mock_user_settings): """Test downgrade_entries when team is not found.""" with patch.dict(os.environ, {'LOCAL_DEPLOYMENT': ''}): with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch( 'storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com' ): with patch.object( LiteLlmManager, '_get_team', new_callable=AsyncMock ) as mock_get_team: mock_get_team.return_value = None result = await LiteLlmManager.downgrade_entries( 'test-org-id', 'test-user-id', mock_user_settings, ) assert result is None @pytest.mark.asyncio async def test_downgrade_entries_successful(self, mock_user_settings): """Test successful downgrade_entries operation.""" mock_response = MagicMock() mock_response.is_success = True mock_response.status_code = 200 mock_response.raise_for_status = MagicMock() mock_team_info_response = MagicMock() mock_team_info_response.is_success = True mock_team_info_response.status_code = 200 mock_team_info_response.json.return_value = { 'team_info': { 'max_budget': 100.0, 'spend': 20.0, }, 'team_memberships': [ { 'user_id': 'test-user-id', 'team_id': 'test-org-id', 'max_budget_in_team': 100.0, 'spend': 20.0, } ], } mock_team_info_response.raise_for_status = MagicMock() mock_key_list_response = MagicMock() mock_key_list_response.is_success = True mock_key_list_response.status_code = 200 mock_key_list_response.json.return_value = { 'keys': ['test-key-1', 'test-key-2'], 'total_count': 2, } mock_key_list_response.raise_for_status = MagicMock() with patch.dict(os.environ, {'LOCAL_DEPLOYMENT': ''}): with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch( 'storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com' ): with patch( 'storage.lite_llm_manager.LITE_LLM_TEAM_ID', 'default-team' ): with patch('httpx.AsyncClient') as mock_client_class: mock_client = AsyncMock() mock_client_class.return_value.__aenter__.return_value = ( mock_client ) # GET requests: get_team (x2 for team info), get_user_keys mock_client.get.side_effect = [ mock_team_info_response, mock_team_info_response, mock_key_list_response, ] mock_client.post.return_value = mock_response result = await LiteLlmManager.downgrade_entries( 'test-org-id', 'test-user-id', mock_user_settings, ) # downgrade_entries returns the user_settings assert result is not None assert ( _agent_value(result.to_settings(), 'agent') == 'TestAgent' ) # Verify downgrade steps were called: # GET requests: # 1. get_team (GET) # 2. get_user_team_info (GET via _get_team) # 3. get_user_keys (GET) # POST requests: # 1. update_user (POST) # 2. add_user_to_team (POST) # 3. update_key for key 1 (POST) # 4. update_key for key 2 (POST) # 5. remove_user_from_team (POST) # 6. delete_team (POST) assert mock_client.get.call_count == 3 assert mock_client.post.call_count == 6 @pytest.mark.asyncio async def test_downgrade_entries_local_deployment(self, mock_user_settings): """Test downgrade_entries in local deployment mode (skips LiteLLM calls).""" with patch.dict(os.environ, {'LOCAL_DEPLOYMENT': 'true'}): with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch( 'storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com' ): result = await LiteLlmManager.downgrade_entries( 'test-org-id', 'test-user-id', mock_user_settings, ) # In local deployment, should return user_settings without # making any LiteLLM calls assert result is not None assert _agent_value(result.to_settings(), 'agent') == 'TestAgent' class TestGetAllKeysForUser: """Test cases for _get_all_keys_for_user method.""" @pytest.mark.asyncio async def test_get_all_keys_missing_config(self): """Test _get_all_keys_for_user when LiteLLM config is missing.""" mock_client = AsyncMock(spec=httpx.AsyncClient) with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', None): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', None): result = await LiteLlmManager._get_all_keys_for_user( mock_client, 'test-user-id' ) assert result == [] mock_client.get.assert_not_called() @pytest.mark.asyncio async def test_get_all_keys_success(self): """Test _get_all_keys_for_user returns keys on success.""" mock_response = MagicMock() mock_response.status_code = 200 mock_response.json.return_value = { 'keys': [ { 'key_name': 'sk-test1234', 'key_alias': 'test-alias', 'team_id': 'test-org', 'metadata': {'type': 'openhands'}, }, { 'key_name': 'sk-test5678', 'key_alias': 'another-alias', 'team_id': 'test-org', 'metadata': None, }, ] } mock_response.raise_for_status = MagicMock() mock_client = AsyncMock(spec=httpx.AsyncClient) mock_client.get.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-api-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): result = await LiteLlmManager._get_all_keys_for_user( mock_client, 'test-user-id' ) assert len(result) == 2 assert result[0]['key_name'] == 'sk-test1234' assert result[1]['key_name'] == 'sk-test5678' # Verify API key header is included mock_client.get.assert_called_once() call_kwargs = mock_client.get.call_args assert call_kwargs.kwargs['headers'] == { 'x-goog-api-key': 'test-api-key' } @pytest.mark.asyncio async def test_get_all_keys_empty_response(self): """Test _get_all_keys_for_user returns empty list when user has no keys.""" mock_response = MagicMock() mock_response.status_code = 200 mock_response.json.return_value = {'keys': []} mock_response.raise_for_status = MagicMock() mock_client = AsyncMock(spec=httpx.AsyncClient) mock_client.get.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-api-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): result = await LiteLlmManager._get_all_keys_for_user( mock_client, 'test-user-id' ) assert result == [] @pytest.mark.asyncio async def test_get_all_keys_api_error(self): """Test _get_all_keys_for_user handles API errors gracefully.""" mock_client = AsyncMock(spec=httpx.AsyncClient) mock_client.get.side_effect = Exception('API Error') with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-api-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): result = await LiteLlmManager._get_all_keys_for_user( mock_client, 'test-user-id' ) assert result == [] class TestVerifyExistingKey: """Test cases for _verify_existing_key method.""" @pytest.mark.asyncio async def test_verify_existing_key_openhands_type_found(self): """Test _verify_existing_key finds matching OpenHands key.""" mock_keys = [ { 'key_name': 'sk-test1234', 'key_alias': 'some-alias', 'team_id': 'test-org', 'metadata': {'type': 'openhands'}, } ] mock_client = AsyncMock(spec=httpx.AsyncClient) with patch.object( LiteLlmManager, '_get_all_keys_for_user', new_callable=AsyncMock ) as mock_get_keys: mock_get_keys.return_value = mock_keys # Key ending with '1234' should match 'sk-test1234' result = await LiteLlmManager._verify_existing_key( mock_client, 'my-key-ending-with-1234', 'test-user-id', 'test-org', openhands_type=True, ) assert result is True @pytest.mark.asyncio async def test_verify_existing_key_openhands_type_not_found(self): """Test _verify_existing_key returns False when key doesn't match.""" mock_keys = [ { 'key_name': 'sk-test1234', 'key_alias': 'some-alias', 'team_id': 'test-org', 'metadata': {'type': 'openhands'}, } ] mock_client = AsyncMock(spec=httpx.AsyncClient) with patch.object( LiteLlmManager, '_get_all_keys_for_user', new_callable=AsyncMock ) as mock_get_keys: mock_get_keys.return_value = mock_keys # Key ending with '5678' should NOT match 'sk-test1234' result = await LiteLlmManager._verify_existing_key( mock_client, 'my-key-ending-with-5678', 'test-user-id', 'test-org', openhands_type=True, ) assert result is False @pytest.mark.asyncio async def test_verify_existing_key_by_alias_openhands_cloud(self): """Test _verify_existing_key finds key by OpenHands Cloud alias.""" user_id = 'test-user-id' org_id = 'test-org' mock_keys = [ { 'key_name': 'sk-testABCD', 'key_alias': get_openhands_cloud_key_alias(user_id, org_id), 'team_id': org_id, 'metadata': None, } ] mock_client = AsyncMock(spec=httpx.AsyncClient) with patch.object( LiteLlmManager, '_get_all_keys_for_user', new_callable=AsyncMock ) as mock_get_keys: mock_get_keys.return_value = mock_keys result = await LiteLlmManager._verify_existing_key( mock_client, 'my-key-ending-with-ABCD', user_id, org_id, openhands_type=False, ) assert result is True @pytest.mark.asyncio async def test_verify_existing_key_by_alias_byor(self): """Test _verify_existing_key finds key by BYOR alias.""" user_id = 'test-user-id' org_id = 'test-org' mock_keys = [ { 'key_name': 'sk-testXYZW', 'key_alias': get_byor_key_alias(user_id, org_id), 'team_id': org_id, 'metadata': None, } ] mock_client = AsyncMock(spec=httpx.AsyncClient) with patch.object( LiteLlmManager, '_get_all_keys_for_user', new_callable=AsyncMock ) as mock_get_keys: mock_get_keys.return_value = mock_keys result = await LiteLlmManager._verify_existing_key( mock_client, 'my-key-ending-with-XYZW', user_id, org_id, openhands_type=False, ) assert result is True @pytest.mark.asyncio async def test_verify_existing_key_wrong_team(self): """Test _verify_existing_key returns False for wrong team_id.""" mock_keys = [ { 'key_name': 'sk-test1234', 'key_alias': 'some-alias', 'team_id': 'different-org', 'metadata': {'type': 'openhands'}, } ] mock_client = AsyncMock(spec=httpx.AsyncClient) with patch.object( LiteLlmManager, '_get_all_keys_for_user', new_callable=AsyncMock ) as mock_get_keys: mock_get_keys.return_value = mock_keys result = await LiteLlmManager._verify_existing_key( mock_client, 'my-key-ending-with-1234', 'test-user-id', 'test-org', openhands_type=True, ) assert result is False @pytest.mark.asyncio async def test_verify_existing_key_no_keys(self): """Test _verify_existing_key returns False when user has no keys.""" mock_client = AsyncMock(spec=httpx.AsyncClient) with patch.object( LiteLlmManager, '_get_all_keys_for_user', new_callable=AsyncMock ) as mock_get_keys: mock_get_keys.return_value = [] result = await LiteLlmManager._verify_existing_key( mock_client, 'some-key-value', 'test-user-id', 'test-org', openhands_type=True, ) assert result is False @pytest.mark.asyncio async def test_verify_existing_key_handles_none_key_name(self): """Test _verify_existing_key handles None key_name gracefully.""" mock_keys = [ { 'key_name': None, 'key_alias': 'some-alias', 'team_id': 'test-org', 'metadata': {'type': 'openhands'}, } ] mock_client = AsyncMock(spec=httpx.AsyncClient) with patch.object( LiteLlmManager, '_get_all_keys_for_user', new_callable=AsyncMock ) as mock_get_keys: mock_get_keys.return_value = mock_keys # Should not raise TypeError, should return False result = await LiteLlmManager._verify_existing_key( mock_client, 'some-key-value', 'test-user-id', 'test-org', openhands_type=True, ) assert result is False @pytest.mark.asyncio async def test_verify_existing_key_handles_empty_key_name(self): """Test _verify_existing_key handles empty key_name gracefully.""" mock_keys = [ { 'key_name': '', 'key_alias': 'some-alias', 'team_id': 'test-org', 'metadata': {'type': 'openhands'}, } ] mock_client = AsyncMock(spec=httpx.AsyncClient) with patch.object( LiteLlmManager, '_get_all_keys_for_user', new_callable=AsyncMock ) as mock_get_keys: mock_get_keys.return_value = mock_keys # Should not raise error, should return False result = await LiteLlmManager._verify_existing_key( mock_client, 'some-key-value', 'test-user-id', 'test-org', openhands_type=True, ) assert result is False class TestBudgetPayloadHandling: """Test cases for budget field handling in API payloads. These tests verify that when max_budget is None, the budget field is NOT included in the JSON payload (which tells LiteLLM to disable budget enforcement), and when max_budget has a value, it IS included. """ @pytest.mark.asyncio async def test_create_team_excludes_max_budget_when_none(self): """Test that _create_team does NOT include max_budget when it is None.""" mock_client = AsyncMock(spec=httpx.AsyncClient) mock_response = MagicMock() mock_response.is_success = True mock_response.status_code = 200 mock_client.post.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-api-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): await LiteLlmManager._create_team( mock_client, team_alias='test-team', team_id='test-team-id', max_budget=None, # None = no budget limit ) # Verify the call was made mock_client.post.assert_called_once() call_args = mock_client.post.call_args # Verify URL assert call_args[0][0] == 'http://test.com/team/new' # Verify that max_budget is NOT in the JSON payload json_payload = call_args[1]['json'] assert 'max_budget' not in json_payload, ( 'max_budget should NOT be in payload when None ' '(omitting it tells LiteLLM to disable budget enforcement)' ) @pytest.mark.asyncio async def test_create_team_includes_max_budget_when_set(self): """Test that _create_team includes max_budget when it has a value.""" mock_client = AsyncMock(spec=httpx.AsyncClient) mock_response = MagicMock() mock_response.is_success = True mock_response.status_code = 200 mock_client.post.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-api-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): await LiteLlmManager._create_team( mock_client, team_alias='test-team', team_id='test-team-id', max_budget=100.0, # Explicit budget limit ) # Verify the call was made mock_client.post.assert_called_once() call_args = mock_client.post.call_args # Verify that max_budget IS in the JSON payload with the correct value json_payload = call_args[1]['json'] assert ( 'max_budget' in json_payload ), 'max_budget should be in payload when set to a value' assert json_payload['max_budget'] == 100.0 @pytest.mark.asyncio async def test_add_user_to_team_excludes_max_budget_when_none(self): """Test that _add_user_to_team does NOT include max_budget_in_team when None.""" mock_client = AsyncMock(spec=httpx.AsyncClient) mock_response = MagicMock() mock_response.is_success = True mock_response.status_code = 200 mock_client.post.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-api-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): await LiteLlmManager._add_user_to_team( mock_client, keycloak_user_id='test-user-id', team_id='test-team-id', max_budget=None, # None = no budget limit ) # Verify the call was made mock_client.post.assert_called_once() call_args = mock_client.post.call_args # Verify URL assert call_args[0][0] == 'http://test.com/team/member_add' # Verify that max_budget_in_team is NOT in the JSON payload json_payload = call_args[1]['json'] assert 'max_budget_in_team' not in json_payload, ( 'max_budget_in_team should NOT be in payload when None ' '(omitting it tells LiteLLM to disable budget enforcement)' ) @pytest.mark.asyncio async def test_add_user_to_team_includes_max_budget_when_set(self): """Test that _add_user_to_team includes max_budget_in_team when set.""" mock_client = AsyncMock(spec=httpx.AsyncClient) mock_response = MagicMock() mock_response.is_success = True mock_response.status_code = 200 mock_client.post.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-api-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): await LiteLlmManager._add_user_to_team( mock_client, keycloak_user_id='test-user-id', team_id='test-team-id', max_budget=50.0, # Explicit budget limit ) # Verify the call was made mock_client.post.assert_called_once() call_args = mock_client.post.call_args # Verify that max_budget_in_team IS in the JSON payload json_payload = call_args[1]['json'] assert ( 'max_budget_in_team' in json_payload ), 'max_budget_in_team should be in payload when set to a value' assert json_payload['max_budget_in_team'] == 50.0 @pytest.mark.asyncio async def test_update_user_in_team_excludes_max_budget_when_none(self): """Test that _update_user_in_team does NOT include max_budget_in_team when None.""" mock_client = AsyncMock(spec=httpx.AsyncClient) mock_response = MagicMock() mock_response.is_success = True mock_response.status_code = 200 mock_client.post.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-api-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): await LiteLlmManager._update_user_in_team( mock_client, keycloak_user_id='test-user-id', team_id='test-team-id', max_budget=None, # None = no budget limit ) # Verify the call was made mock_client.post.assert_called_once() call_args = mock_client.post.call_args # Verify URL assert call_args[0][0] == 'http://test.com/team/member_update' # Verify that max_budget_in_team is NOT in the JSON payload json_payload = call_args[1]['json'] assert 'max_budget_in_team' not in json_payload, ( 'max_budget_in_team should NOT be in payload when None ' '(omitting it tells LiteLLM to disable budget enforcement)' ) @pytest.mark.asyncio async def test_update_user_in_team_includes_max_budget_when_set(self): """Test that _update_user_in_team includes max_budget_in_team when set.""" mock_client = AsyncMock(spec=httpx.AsyncClient) mock_response = MagicMock() mock_response.is_success = True mock_response.status_code = 200 mock_client.post.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-api-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): await LiteLlmManager._update_user_in_team( mock_client, keycloak_user_id='test-user-id', team_id='test-team-id', max_budget=75.0, # Explicit budget limit ) # Verify the call was made mock_client.post.assert_called_once() call_args = mock_client.post.call_args # Verify that max_budget_in_team IS in the JSON payload json_payload = call_args[1]['json'] assert ( 'max_budget_in_team' in json_payload ), 'max_budget_in_team should be in payload when set to a value' assert json_payload['max_budget_in_team'] == 75.0 class TestGetTeamMembersFinancialData: """Test cases for _get_team_members_financial_data method.""" @pytest.fixture def mock_http_client(self): """Create a mock HTTP client.""" return AsyncMock(spec=httpx.AsyncClient) @pytest.mark.asyncio async def test_returns_financial_data_for_all_team_members(self, mock_http_client): """ GIVEN: Team with multiple members having financial data WHEN: _get_team_members_financial_data is called THEN: Returns dict with team info and member data """ # Arrange mock_response = MagicMock() mock_response.is_success = True mock_response.status_code = 200 mock_response.json.return_value = { 'team_info': {'team_id': 'test-team', 'max_budget': 500.0, 'spend': 125.5}, 'team_memberships': [ { 'user_id': 'user-1', 'spend': 50.0, 'max_budget_in_team': 200.0, }, { 'user_id': 'user-2', 'spend': 75.5, 'max_budget_in_team': 150.0, }, ], } mock_response.raise_for_status = MagicMock() mock_http_client.get.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): # Act result = await LiteLlmManager._get_team_members_financial_data( mock_http_client, 'test-team' ) # Assert assert result['team_max_budget'] == 500.0 assert result['team_spend'] == 125.5 assert len(result['members']) == 2 # Both users have individual budgets (max_budget_in_team is set) assert result['members']['user-1'] == { 'spend': 50.0, 'max_budget': 200.0, 'uses_shared_budget': False, } assert result['members']['user-2'] == { 'spend': 75.5, 'max_budget': 150.0, 'uses_shared_budget': False, } @pytest.mark.asyncio async def test_returns_empty_dict_when_litellm_not_configured( self, mock_http_client ): """ GIVEN: LiteLLM API key or URL not configured WHEN: _get_team_members_financial_data is called THEN: Returns empty dict """ # Arrange - no patching, so LITE_LLM_API_KEY/URL are None with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', None): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', None): # Act result = await LiteLlmManager._get_team_members_financial_data( mock_http_client, 'test-team' ) # Assert assert result == {} mock_http_client.get.assert_not_called() @pytest.mark.asyncio async def test_returns_empty_dict_when_team_not_found(self, mock_http_client): """ GIVEN: Team does not exist in LiteLLM WHEN: _get_team_members_financial_data is called THEN: Returns empty dict """ # Arrange mock_response = MagicMock() mock_response.status_code = 404 mock_response.raise_for_status.side_effect = httpx.HTTPStatusError( 'Not found', request=MagicMock(), response=mock_response ) mock_http_client.get.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): # Act & Assert with pytest.raises(httpx.HTTPStatusError): await LiteLlmManager._get_team_members_financial_data( mock_http_client, 'nonexistent-team' ) @pytest.mark.asyncio async def test_returns_empty_members_when_team_has_no_members( self, mock_http_client ): """ GIVEN: Team exists but has no members WHEN: _get_team_members_financial_data is called THEN: Returns structure with empty members dict """ # Arrange mock_response = MagicMock() mock_response.is_success = True mock_response.status_code = 200 mock_response.json.return_value = { 'team_info': {'team_id': 'empty-team', 'max_budget': 100.0, 'spend': 0}, 'team_memberships': [], } mock_response.raise_for_status = MagicMock() mock_http_client.get.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): # Act result = await LiteLlmManager._get_team_members_financial_data( mock_http_client, 'empty-team' ) # Assert assert result['team_max_budget'] == 100.0 assert result['team_spend'] == 0 assert result['members'] == {} @pytest.mark.asyncio async def test_falls_back_to_team_budget_when_member_budget_missing( self, mock_http_client ): """ GIVEN: Team with shared budget, members without individual max_budget_in_team WHEN: _get_team_members_financial_data is called THEN: Falls back to team_info.max_budget for members without individual budget """ # Arrange mock_response = MagicMock() mock_response.is_success = True mock_response.status_code = 200 mock_response.json.return_value = { 'team_info': {'team_id': 'test-team', 'max_budget': 500.0, 'spend': 150.0}, 'team_memberships': [ { 'user_id': 'user-no-individual-budget', 'spend': 50.0, # No max_budget_in_team - should fall back to team budget }, { 'user_id': 'user-with-individual-budget', 'spend': 75.0, 'max_budget_in_team': 200.0, # Individual budget set }, { 'user_id': 'user-null-budget', 'spend': 25.0, 'max_budget_in_team': None, # Explicit null - fall back to team }, ], } mock_response.raise_for_status = MagicMock() mock_http_client.get.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): # Act result = await LiteLlmManager._get_team_members_financial_data( mock_http_client, 'test-team' ) # Assert assert result['team_max_budget'] == 500.0 assert result['team_spend'] == 150.0 members = result['members'] assert members['user-no-individual-budget'] == { 'spend': 50.0, 'max_budget': 500.0, 'uses_shared_budget': True, } assert members['user-with-individual-budget'] == { 'spend': 75.0, 'max_budget': 200.0, 'uses_shared_budget': False, } assert members['user-null-budget'] == { 'spend': 25.0, 'max_budget': 500.0, 'uses_shared_budget': True, } @pytest.mark.asyncio async def test_uses_defaults_when_no_budget_data_available(self, mock_http_client): """ GIVEN: Team without budget and members without individual budgets WHEN: _get_team_members_financial_data is called THEN: Returns default values (spend=0, max_budget=None) """ # Arrange mock_response = MagicMock() mock_response.is_success = True mock_response.status_code = 200 mock_response.json.return_value = { 'team_info': {'team_id': 'test-team'}, # No max_budget at team level 'team_memberships': [ { 'user_id': 'user-no-data', # No spend or max_budget_in_team }, { 'user_id': 'user-null-spend', 'spend': None, # Explicit null }, ], } mock_response.raise_for_status = MagicMock() mock_http_client.get.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): # Act result = await LiteLlmManager._get_team_members_financial_data( mock_http_client, 'test-team' ) # Assert assert result['team_max_budget'] is None assert result['team_spend'] == 0 members = result['members'] # Both users fall back to team budget (which is None) assert members['user-no-data'] == { 'spend': 0, 'max_budget': None, 'uses_shared_budget': True, } assert members['user-null-spend'] == { 'spend': 0, 'max_budget': None, 'uses_shared_budget': True, } @pytest.mark.asyncio async def test_skips_members_without_user_id(self, mock_http_client): """ GIVEN: Team with members, some missing user_id WHEN: _get_team_members_financial_data is called THEN: Skips members without user_id """ # Arrange mock_response = MagicMock() mock_response.is_success = True mock_response.status_code = 200 mock_response.json.return_value = { 'team_info': {'team_id': 'test-team', 'max_budget': 300.0, 'spend': 105.0}, 'team_memberships': [ { 'user_id': 'valid-user', 'spend': 25.0, 'max_budget_in_team': 100.0, }, { # Missing user_id 'spend': 50.0, 'max_budget_in_team': 200.0, }, { 'user_id': None, # Explicit null 'spend': 30.0, }, ], } mock_response.raise_for_status = MagicMock() mock_http_client.get.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): # Act result = await LiteLlmManager._get_team_members_financial_data( mock_http_client, 'test-team' ) # Assert - only valid user should be included assert result['team_max_budget'] == 300.0 assert result['team_spend'] == 105.0 assert len(result['members']) == 1 assert 'valid-user' in result['members'] assert result['members']['valid-user'] == { 'spend': 25.0, 'max_budget': 100.0, 'uses_shared_budget': False, }