Add default initial budget for teams/users (#13389)

Co-authored-by: OpenHands Bot <contact@all-hands.dev>
Co-authored-by: openhands <openhands@all-hands.dev>

enterprise/tests/unit/test_lite_llm_manager.py

@@ -2,7 +2,9 @@
 Unit tests for LiteLlmManager class.
 """
 
+import importlib
 import os
+import sys
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
@@ -21,6 +23,71 @@ from storage.user_settings import UserSettings
 from openhands.server.settings import Settings
 
 
+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 var
+        os.environ.pop('DEFAULT_INITIAL_BUDGET', 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_defaults_to_zero(self):
+        """Test that DEFAULT_INITIAL_BUDGET defaults to 0.0 when env var not set."""
+        # 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']
+
+        # Clear the env var and reimport
+        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(self):
+        """Test that DEFAULT_INITIAL_BUDGET uses value from environment variable."""
+        if 'storage.lite_llm_manager' in sys.modules:
+            del sys.modules['storage.lite_llm_manager']
+
+        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_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['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['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."""
 
@@ -242,10 +309,10 @@ class TestLiteLlmManager:
             assert add_user_call[1]['json']['max_budget_in_team'] == 30.0
 
     @pytest.mark.asyncio
-    async def test_create_entries_new_org_uses_zero_budget(
+    async def test_create_entries_new_org_uses_default_initial_budget(
         self, mock_settings, mock_response
     ):
-        """Test that create_entries uses budget=0 for new org (team doesn't exist)."""
+        """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
@@ -273,6 +340,7 @@ class TestLiteLlmManager:
             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
@@ -280,16 +348,67 @@ class TestLiteLlmManager:
 
             assert result is not None
 
-            # Verify _create_team was called with budget=0
+            # 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 budget=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_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_404_response
+        mock_client.get.return_value.raise_for_status.side_effect = (
+            httpx.HTTPStatusError(
+                message='Not Found', request=MagicMock(), response=mock_404_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."""