OpenHands/enterprise/tests/unit/test_saas_settings_store.py

from unittest.mock import AsyncMock, MagicMock, patch

import httpx
import pytest
from pydantic import SecretStr
from server.constants import (
    CURRENT_USER_SETTINGS_VERSION,
    LITE_LLM_API_URL,
    LITE_LLM_TEAM_ID,
    get_default_litellm_model,
)
from storage.saas_settings_store import SaasSettingsStore
from storage.user_settings import UserSettings

from openhands.core.config.openhands_config import OpenHandsConfig
from openhands.server.settings import Settings


@pytest.fixture
def mock_litellm_get_response():
    mock_response = AsyncMock()
    mock_response.is_success = True
    mock_response.json = MagicMock(return_value={'user_info': {}})
    return mock_response


@pytest.fixture
def mock_litellm_post_response():
    mock_response = AsyncMock()
    mock_response.is_success = True
    mock_response.json = MagicMock(return_value={'key': 'test_api_key'})
    return mock_response


@pytest.fixture
def mock_litellm_api(mock_litellm_get_response, mock_litellm_post_response):
    api_key_patch = patch('storage.saas_settings_store.LITE_LLM_API_KEY', 'test_key')
    api_url_patch = patch(
        'storage.saas_settings_store.LITE_LLM_API_URL', 'http://test.url'
    )
    team_id_patch = patch('storage.saas_settings_store.LITE_LLM_TEAM_ID', 'test_team')
    client_patch = patch('httpx.AsyncClient')

    with api_key_patch, api_url_patch, team_id_patch, client_patch as mock_client:
        mock_client.return_value.__aenter__.return_value.get.return_value = (
            mock_litellm_get_response
        )
        mock_client.return_value.__aenter__.return_value.post.return_value = (
            mock_litellm_post_response
        )
        yield mock_client


@pytest.fixture
def mock_stripe():
    search_patch = patch(
        'stripe.Customer.search_async',
        AsyncMock(return_value=MagicMock(id='mock-customer-id')),
    )
    payment_patch = patch(
        'stripe.Customer.list_payment_methods_async',
        AsyncMock(return_value=MagicMock(data=[{}])),
    )
    with search_patch, payment_patch:
        yield


@pytest.fixture
def mock_github_user():
    with patch(
        'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
        AsyncMock(return_value={'attributes': {'github_id': ['12345']}}),
    ) as mock_github:
        yield mock_github


@pytest.fixture
def mock_config():
    config = MagicMock(spec=OpenHandsConfig)
    config.jwt_secret = SecretStr('test_secret')
    config.file_store = 'google_cloud'
    config.file_store_path = 'bucket'
    return config


@pytest.fixture
def settings_store(session_maker, mock_config):
    store = SaasSettingsStore('user-id', session_maker, mock_config)

    # Patch the store method directly to filter out email and email_verified
    original_load = store.load
    original_create_default = store.create_default_settings
    original_update_litellm = store.update_settings_with_litellm_default

    # Patch the load method to add email and email_verified
    async def patched_load():
        settings = await original_load()
        if settings:
            # Add email and email_verified fields to mimic SaasUserAuth behavior
            settings.email = 'test@example.com'
            settings.email_verified = True
        return settings

    # Patch the create_default_settings method to add email and email_verified
    async def patched_create_default(settings):
        settings = await original_create_default(settings)
        if settings:
            # Add email and email_verified fields to mimic SaasUserAuth behavior
            settings.email = 'test@example.com'
            settings.email_verified = True
        return settings

    # Patch the update_settings_with_litellm_default method
    async def patched_update_litellm(settings):
        updated_settings = await original_update_litellm(settings)
        if updated_settings:
            # Add email and email_verified fields to mimic SaasUserAuth behavior
            updated_settings.email = 'test@example.com'
            updated_settings.email_verified = True
        return updated_settings

    # Patch the store method to filter out email and email_verified
    async def patched_store(item):
        if item:
            # Make a copy of the item without email and email_verified
            item_dict = item.model_dump(context={'expose_secrets': True})
            if 'email' in item_dict:
                del item_dict['email']
            if 'email_verified' in item_dict:
                del item_dict['email_verified']
            if 'secrets_store' in item_dict:
                del item_dict['secrets_store']

            # Continue with the original implementation
            with store.session_maker() as session:
                existing = None
                if item_dict:
                    store._encrypt_kwargs(item_dict)
                    query = session.query(UserSettings).filter(
                        UserSettings.keycloak_user_id == store.user_id
                    )

                    # First check if we have an existing entry in the new table
                    existing = query.first()

                if existing:
                    # Update existing entry
                    for key, value in item_dict.items():
                        if key in existing.__class__.__table__.columns:
                            setattr(existing, key, value)
                    existing.user_version = CURRENT_USER_SETTINGS_VERSION
                    session.merge(existing)
                else:
                    item_dict['keycloak_user_id'] = store.user_id
                    item_dict['user_version'] = CURRENT_USER_SETTINGS_VERSION
                    settings = UserSettings(**item_dict)
                    session.add(settings)
                session.commit()

    # Replace the methods with our patched versions
    store.store = patched_store
    store.load = patched_load
    store.create_default_settings = patched_create_default
    store.update_settings_with_litellm_default = patched_update_litellm
    return store


@pytest.mark.asyncio
async def test_store_and_load_keycloak_user(settings_store):
    # Set a UUID-like Keycloak user ID
    settings_store.user_id = '550e8400-e29b-41d4-a716-446655440000'
    settings = Settings(
        llm_api_key=SecretStr('secret_key'),
        llm_base_url=LITE_LLM_API_URL,
        agent='smith',
        email='test@example.com',
        email_verified=True,
    )

    await settings_store.store(settings)

    # Load and verify settings
    loaded_settings = await settings_store.load()
    assert loaded_settings is not None
    assert loaded_settings.llm_api_key.get_secret_value() == 'secret_key'
    assert loaded_settings.agent == 'smith'

    # Verify it was stored in user_settings table with keycloak_user_id
    with settings_store.session_maker() as session:
        stored = (
            session.query(UserSettings)
            .filter(
                UserSettings.keycloak_user_id == '550e8400-e29b-41d4-a716-446655440000'
            )
            .first()
        )
        assert stored is not None
        assert stored.agent == 'smith'


@pytest.mark.asyncio
async def test_load_returns_default_when_not_found(
    settings_store, mock_litellm_api, mock_stripe, mock_github_user, session_maker
):
    file_store = MagicMock()
    file_store.read.side_effect = FileNotFoundError()

    with (
        patch(
            'storage.saas_settings_store.get_file_store',
            MagicMock(return_value=file_store),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        loaded_settings = await settings_store.load()
        assert loaded_settings is not None
        assert loaded_settings.language == 'en'
        assert loaded_settings.agent == 'CodeActAgent'
        assert loaded_settings.llm_api_key.get_secret_value() == 'test_api_key'
        assert loaded_settings.llm_base_url == 'http://test.url'


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default(
    settings_store, mock_litellm_api, session_maker
):
    settings = Settings()
    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'testy@tester.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        settings = await settings_store.update_settings_with_litellm_default(settings)

    assert settings.agent == 'CodeActAgent'
    assert settings.llm_api_key
    assert settings.llm_api_key.get_secret_value() == 'test_api_key'
    assert settings.llm_base_url == 'http://test.url'

    # Get the actual call arguments
    call_args = mock_litellm_api.return_value.__aenter__.return_value.post.call_args[1]

    # Check that the URL and most of the JSON payload match what we expect
    assert call_args['json']['user_email'] == 'testy@tester.com'
    assert call_args['json']['models'] == []
    assert call_args['json']['max_budget'] == 10.0
    assert call_args['json']['user_id'] == 'user-id'
    assert call_args['json']['teams'] == ['test_team']
    assert call_args['json']['auto_create_key'] is True
    assert call_args['json']['send_invite_email'] is False
    assert call_args['json']['metadata']['version'] == CURRENT_USER_SETTINGS_VERSION
    assert 'model' in call_args['json']['metadata']


@pytest.mark.asyncio
async def test_create_default_settings_no_user_id():
    store = SaasSettingsStore('', MagicMock(), MagicMock())
    settings = await store.create_default_settings(None)
    assert settings is None


@pytest.mark.asyncio
async def test_create_default_settings_require_payment_enabled(
    settings_store, mock_stripe
):
    # Mock stripe_service.has_payment_method to return False
    with (
        patch('storage.saas_settings_store.REQUIRE_PAYMENT', True),
        patch(
            'stripe.Customer.list_payment_methods_async',
            AsyncMock(return_value=MagicMock(data=[])),
        ),
        patch(
            'integrations.stripe_service.session_maker', settings_store.session_maker
        ),
    ):
        settings = await settings_store.create_default_settings(None)
        assert settings is None


@pytest.mark.asyncio
async def test_create_default_settings_require_payment_disabled(
    settings_store, mock_stripe, mock_github_user, mock_litellm_api, session_maker
):
    # Even without payment method, should get default settings when REQUIRE_PAYMENT is False
    file_store = MagicMock()
    file_store.read.side_effect = FileNotFoundError()
    with (
        patch('storage.saas_settings_store.REQUIRE_PAYMENT', False),
        patch(
            'stripe.Customer.list_payment_methods_async',
            AsyncMock(return_value=MagicMock(data=[])),
        ),
        patch(
            'storage.saas_settings_store.get_file_store',
            MagicMock(return_value=file_store),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        settings = await settings_store.create_default_settings(None)
        assert settings is not None
        assert settings.language == 'en'


@pytest.mark.asyncio
async def test_create_default_lite_llm_settings_no_api_config(settings_store):
    with (
        patch('storage.saas_settings_store.LITE_LLM_API_KEY', None),
        patch('storage.saas_settings_store.LITE_LLM_API_URL', None),
    ):
        settings = Settings()
        settings = await settings_store.update_settings_with_litellm_default(settings)


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default_error(settings_store):
    with patch(
        'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
        AsyncMock(return_value={'email': 'duplicate@example.com'}),
    ):
        with patch('httpx.AsyncClient') as mock_client:
            mock_client.return_value.__aenter__.return_value.get.return_value = (
                AsyncMock(
                    json=MagicMock(
                        return_value={'user_info': {'max_budget': 10, 'spend': 5}}
                    )
                )
            )
            mock_client.return_value.__aenter__.return_value.post.return_value.is_success = False
            settings = Settings()
            settings = await settings_store.update_settings_with_litellm_default(
                settings
            )
            assert settings is None


@pytest.mark.asyncio
@pytest.mark.parametrize(
    'status_code,user_info_response,should_succeed',
    [
        # 200 OK with user info - existing user (v1.79.x and v1.80+ behavior)
        (200, {'user_info': {'max_budget': 10, 'spend': 5}}, True),
        # 200 OK with empty user info - new user (v1.79.x behavior)
        (200, {'user_info': None}, True),
        # 404 Not Found - new user (v1.80+ behavior)
        (404, None, True),
        # 500 Internal Server Error - should fail
        (500, None, False),
    ],
)
async def test_update_settings_with_litellm_default_handles_user_info_responses(
    settings_store, session_maker, status_code, user_info_response, should_succeed
):
    """Test that various LiteLLM user/info responses are handled correctly.

    LiteLLM API behavior changed between versions:
    - v1.79.x and earlier: GET /user/info always succeeds with empty user_info
    - v1.80.x and later: GET /user/info returns 404 for non-existent users
    """
    mock_get_response = MagicMock()
    mock_get_response.status_code = status_code
    if user_info_response is not None:
        mock_get_response.json = MagicMock(return_value=user_info_response)
        mock_get_response.raise_for_status = MagicMock()
    else:
        mock_get_response.raise_for_status = MagicMock(
            side_effect=httpx.HTTPStatusError(
                'Error', request=MagicMock(), response=mock_get_response
            )
            if status_code >= 500
            else None
        )

    # Mock successful responses for POST operations (delete and create)
    mock_post_response = MagicMock()
    mock_post_response.is_success = True
    mock_post_response.json = MagicMock(return_value={'key': 'new_user_api_key'})

    with (
        patch('storage.saas_settings_store.LITE_LLM_API_KEY', 'test_key'),
        patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://test.url'),
        patch('storage.saas_settings_store.LITE_LLM_TEAM_ID', 'test_team'),
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'testuser@example.com'}),
        ),
        patch('httpx.AsyncClient') as mock_client,
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        # Set up the mock client
        mock_client.return_value.__aenter__.return_value.get.return_value = (
            mock_get_response
        )
        mock_client.return_value.__aenter__.return_value.post.return_value = (
            mock_post_response
        )

        settings = Settings()
        if should_succeed:
            settings = await settings_store.update_settings_with_litellm_default(
                settings
            )
            assert settings is not None
            assert settings.llm_api_key is not None
            assert settings.llm_api_key.get_secret_value() == 'new_user_api_key'
        else:
            with pytest.raises(httpx.HTTPStatusError):
                await settings_store.update_settings_with_litellm_default(settings)


@pytest.mark.asyncio
async def test_update_settings_with_litellm_retry_on_duplicate_email(
    settings_store, mock_litellm_api, session_maker
):
    # First response is a delete and succeeds
    mock_delete_response = MagicMock()
    mock_delete_response.is_success = True
    mock_delete_response.status_code = 200

    # Second response fails with duplicate email error
    mock_error_response = MagicMock()
    mock_error_response.is_success = False
    mock_error_response.status_code = 400
    mock_error_response.text = 'User with this email already exists'

    # Thire response succeeds with no email
    mock_success_response = MagicMock()
    mock_success_response.is_success = True
    mock_success_response.json = MagicMock(return_value={'key': 'new_test_api_key'})

    # Set up mocks
    post_mock = AsyncMock()
    post_mock.side_effect = [
        mock_delete_response,
        mock_error_response,
        mock_success_response,
    ]
    mock_litellm_api.return_value.__aenter__.return_value.post = post_mock

    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'duplicate@example.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        settings = Settings()
        settings = await settings_store.update_settings_with_litellm_default(settings)

    assert settings is not None
    assert settings.llm_api_key
    assert settings.llm_api_key.get_secret_value() == 'new_test_api_key'

    # Verify second call was with email
    second_call_args = post_mock.call_args_list[1][1]
    assert second_call_args['json']['user_email'] == 'duplicate@example.com'

    # Verify third call was with None for email
    third_call_args = post_mock.call_args_list[2][1]
    assert third_call_args['json']['user_email'] is None


@pytest.mark.asyncio
async def test_create_user_in_lite_llm(settings_store):
    # Test the _create_user_in_lite_llm method directly
    mock_client = AsyncMock()
    mock_response = AsyncMock()
    mock_response.is_success = True
    mock_client.post.return_value = mock_response
    test_model = 'custom-model/test-model'

    # Test with email
    await settings_store._create_user_in_lite_llm(
        mock_client, 'test@example.com', 50, 10, test_model
    )

    # Get the actual call arguments
    call_args = mock_client.post.call_args[1]

    # Check that the URL and most of the JSON payload match what we expect
    assert call_args['json']['user_email'] == 'test@example.com'
    assert call_args['json']['models'] == []
    assert call_args['json']['max_budget'] == 50
    assert call_args['json']['spend'] == 10
    assert call_args['json']['user_id'] == 'user-id'
    assert call_args['json']['teams'] == [LITE_LLM_TEAM_ID]
    assert call_args['json']['auto_create_key'] is True
    assert call_args['json']['send_invite_email'] is False
    assert call_args['json']['metadata']['version'] == CURRENT_USER_SETTINGS_VERSION
    assert call_args['json']['metadata']['model'] == test_model

    # Test with None email
    mock_client.post.reset_mock()
    await settings_store._create_user_in_lite_llm(mock_client, None, 25, 15, test_model)

    # Get the actual call arguments
    call_args = mock_client.post.call_args[1]

    # Check that the URL and most of the JSON payload match what we expect
    assert call_args['json']['user_email'] is None
    assert call_args['json']['models'] == []
    assert call_args['json']['max_budget'] == 25
    assert call_args['json']['spend'] == 15
    assert call_args['json']['user_id'] == str(settings_store.user_id)
    assert call_args['json']['teams'] == [LITE_LLM_TEAM_ID]
    assert call_args['json']['auto_create_key'] is True
    assert call_args['json']['send_invite_email'] is False
    assert call_args['json']['metadata']['version'] == CURRENT_USER_SETTINGS_VERSION
    assert call_args['json']['metadata']['model'] == test_model

    # Verify response is returned correctly
    assert (
        await settings_store._create_user_in_lite_llm(
            mock_client, 'email@test.com', 30, 7, test_model
        )
        == mock_response
    )


@pytest.mark.asyncio
async def test_encryption(settings_store):
    settings_store.user_id = 'mock-id'  # GitHub user ID
    settings = Settings(
        llm_api_key=SecretStr('secret_key'),
        agent='smith',
        llm_base_url=LITE_LLM_API_URL,
        email='test@example.com',
        email_verified=True,
    )
    await settings_store.store(settings)
    with settings_store.session_maker() as session:
        stored = (
            session.query(UserSettings)
            .filter(UserSettings.keycloak_user_id == 'mock-id')
            .first()
        )
        # The stored key should be encrypted
        assert stored.llm_api_key != 'secret_key'
        # But we should be able to decrypt it when loading
        loaded_settings = await settings_store.load()
        assert loaded_settings.llm_api_key.get_secret_value() == 'secret_key'


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default_preserves_custom_model(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User has a custom LLM model set
    custom_model = 'anthropic/claude-3-5-sonnet-20241022'
    settings = Settings(llm_model=custom_model)

    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        # Act: Update settings with LiteLLM defaults
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

    # Assert: Custom model is preserved
    assert updated_settings is not None
    assert updated_settings.llm_model == custom_model
    assert updated_settings.agent == 'CodeActAgent'
    assert updated_settings.llm_api_key is not None

    # Assert: LiteLLM metadata contains user's custom model
    call_args = mock_litellm_api.return_value.__aenter__.return_value.post.call_args[1]
    assert call_args['json']['metadata']['model'] == custom_model


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default_uses_default_when_no_model(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User has no model set (new user scenario)
    settings = Settings()

    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'newuser@example.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        # Act: Update settings with LiteLLM defaults
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

    # Assert: Default model is assigned
    assert updated_settings is not None
    expected_default = get_default_litellm_model()
    assert updated_settings.llm_model == expected_default
    assert updated_settings.agent == 'CodeActAgent'

    # Assert: LiteLLM metadata contains default model
    call_args = mock_litellm_api.return_value.__aenter__.return_value.post.call_args[1]
    assert call_args['json']['metadata']['model'] == expected_default


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default_handles_empty_string_model(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User has empty string as model (edge case)
    settings = Settings(llm_model='')

    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        # Act: Update settings with LiteLLM defaults
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

    # Assert: Default model is used (empty string treated as no model)
    assert updated_settings is not None
    expected_default = get_default_litellm_model()
    assert updated_settings.llm_model == expected_default


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default_handles_whitespace_model(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User has whitespace-only model (edge case)
    settings = Settings(llm_model='   ')

    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        # Act: Update settings with LiteLLM defaults
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

    # Assert: Default model is used (whitespace treated as no model)
    assert updated_settings is not None
    expected_default = get_default_litellm_model()
    assert updated_settings.llm_model == expected_default


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default_preserves_custom_api_key(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User has a custom API key and custom model (so has_custom=True)
    custom_api_key = 'sk-custom-user-api-key-12345'
    custom_model = 'gpt-4'
    settings = Settings(llm_model=custom_model, llm_api_key=SecretStr(custom_api_key))

    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        # Act: Update settings with LiteLLM defaults
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

    # Assert: Custom API key is preserved when user has custom settings
    assert updated_settings is not None
    assert updated_settings.llm_api_key.get_secret_value() == custom_api_key
    assert updated_settings.llm_api_key.get_secret_value() != 'test_api_key'


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default_preserves_custom_base_url(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User has a custom base URL
    custom_base_url = 'https://api.custom-llm-provider.com/v1'
    settings = Settings(llm_base_url=custom_base_url)

    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        # Act: Update settings with LiteLLM defaults
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

    # Assert: Custom base URL is preserved
    assert updated_settings is not None
    assert updated_settings.llm_base_url == custom_base_url
    assert updated_settings.llm_base_url != LITE_LLM_API_URL


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default_preserves_custom_api_key_and_base_url(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User has both custom API key and base URL
    custom_api_key = 'sk-custom-user-api-key-67890'
    custom_base_url = 'https://api.another-llm-provider.com/v1'
    custom_model = 'openai/gpt-4'
    settings = Settings(
        llm_model=custom_model,
        llm_api_key=SecretStr(custom_api_key),
        llm_base_url=custom_base_url,
    )

    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        # Act: Update settings with LiteLLM defaults
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

    # Assert: All custom settings are preserved
    assert updated_settings is not None
    assert updated_settings.llm_model == custom_model
    assert updated_settings.llm_api_key.get_secret_value() == custom_api_key
    assert updated_settings.llm_base_url == custom_base_url


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default_uses_default_api_key_when_none(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User has no API key set
    settings = Settings(llm_api_key=None)

    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        # Act: Update settings with LiteLLM defaults
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

    # Assert: Default LiteLLM API key is assigned
    assert updated_settings is not None
    assert updated_settings.llm_api_key is not None
    assert updated_settings.llm_api_key.get_secret_value() == 'test_api_key'


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default_uses_default_base_url_when_none(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User has no base URL set
    settings = Settings(llm_base_url=None)

    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
        patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://test.url'),
    ):
        # Act: Update settings with LiteLLM defaults
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

    # Assert: Default LiteLLM base URL is assigned (using mocked value)
    assert updated_settings is not None
    assert updated_settings.llm_base_url == 'http://test.url'


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default_handles_empty_api_key(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User has empty string as API key (edge case)
    settings = Settings(llm_api_key=SecretStr(''))

    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        # Act: Update settings with LiteLLM defaults
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

    # Assert: Default API key is used (empty string treated as no key)
    assert updated_settings is not None
    assert updated_settings.llm_api_key.get_secret_value() == 'test_api_key'


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default_handles_empty_base_url(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User has empty string as base URL (edge case)
    settings = Settings(llm_base_url='')

    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
        patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://test.url'),
    ):
        # Act: Update settings with LiteLLM defaults
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

    # Assert: Default base URL is used (empty string treated as no URL)
    assert updated_settings is not None
    assert updated_settings.llm_base_url == 'http://test.url'


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default_handles_whitespace_api_key(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User has whitespace-only API key (edge case)
    settings = Settings(llm_api_key=SecretStr('   '))

    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
    ):
        # Act: Update settings with LiteLLM defaults
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

    # Assert: Default API key is used (whitespace treated as no key)
    assert updated_settings is not None
    assert updated_settings.llm_api_key.get_secret_value() == 'test_api_key'


@pytest.mark.asyncio
async def test_update_settings_with_litellm_default_handles_whitespace_base_url(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User has whitespace-only base URL (edge case)
    settings = Settings(llm_base_url='   ')

    with (
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
        patch('storage.saas_settings_store.session_maker', session_maker),
        patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://test.url'),
    ):
        # Act: Update settings with LiteLLM defaults
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

    # Assert: Default base URL is used (whitespace treated as no URL)
    assert updated_settings is not None
    assert updated_settings.llm_base_url == 'http://test.url'


# Tests for version migration and helper methods


@pytest.mark.asyncio
async def test_has_custom_settings_with_custom_base_url(settings_store):
    # Arrange: User with custom base URL (BYOR)
    with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
        settings = Settings(llm_base_url='http://custom.url')

        # Act: Check if has custom settings
        has_custom = settings_store._has_custom_settings(settings, None)

        # Assert: Custom base URL detected
        assert has_custom is True


@pytest.mark.asyncio
async def test_has_custom_settings_with_default_base_url(settings_store):
    # Arrange: User with default base URL
    with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
        settings = Settings(llm_base_url='http://default.url')

        # Act: Check if has custom settings
        has_custom = settings_store._has_custom_settings(settings, None)

        # Assert: No custom settings (no model set)
        assert has_custom is False


@pytest.mark.asyncio
async def test_has_custom_settings_with_no_model(settings_store):
    # Arrange: User with no model set
    with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
        settings = Settings(llm_model=None, llm_base_url='http://default.url')

        # Act: Check if has custom settings
        has_custom = settings_store._has_custom_settings(settings, None)

        # Assert: No custom settings (using defaults)
        assert has_custom is False


@pytest.mark.asyncio
async def test_has_custom_settings_with_empty_model(settings_store):
    # Arrange: User with empty model
    with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
        settings = Settings(llm_model='', llm_base_url='http://default.url')

        # Act: Check if has custom settings
        has_custom = settings_store._has_custom_settings(settings, None)

        # Assert: No custom settings (empty treated as no model)
        assert has_custom is False


@pytest.mark.asyncio
async def test_has_custom_settings_with_whitespace_model(settings_store):
    # Arrange: User with whitespace-only model
    with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
        settings = Settings(llm_model='   ', llm_base_url='http://default.url')

        # Act: Check if has custom settings
        has_custom = settings_store._has_custom_settings(settings, None)

        # Assert: No custom settings (whitespace treated as no model)
        assert has_custom is False


@pytest.mark.asyncio
async def test_has_custom_settings_with_custom_model(settings_store):
    # Arrange: User with custom model
    with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
        settings = Settings(llm_model='gpt-4', llm_base_url='http://default.url')

        # Act: Check if has custom settings
        has_custom = settings_store._has_custom_settings(settings, None)

        # Assert: Custom model detected
        assert has_custom is True


@pytest.mark.asyncio
async def test_has_custom_settings_matches_old_default_model(settings_store):
    # Arrange: User with old version and model matching old default
    with (
        patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'),
        patch('server.constants.CURRENT_USER_SETTINGS_VERSION', 5),
        patch(
            'server.constants.USER_SETTINGS_VERSION_TO_MODEL',
            {1: 'claude-3-5-sonnet-20241022'},
        ),
    ):
        settings = Settings(
            llm_model='litellm_proxy/prod/claude-3-5-sonnet-20241022',
            llm_base_url='http://default.url',
        )

        # Act: Check if has custom settings
        has_custom = settings_store._has_custom_settings(settings, 1)

        # Assert: Matches old default, so not custom
        assert has_custom is False


@pytest.mark.asyncio
async def test_has_custom_settings_matches_old_default_by_base_name(settings_store):
    # Arrange: User with old version and model matching old default by base name
    with (
        patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'),
        patch('server.constants.CURRENT_USER_SETTINGS_VERSION', 5),
        patch(
            'server.constants.USER_SETTINGS_VERSION_TO_MODEL',
            {1: 'claude-3-5-sonnet-20241022'},
        ),
    ):
        settings = Settings(
            llm_model='anthropic/claude-3-5-sonnet-20241022',
            llm_base_url='http://default.url',
        )

        # Act: Check if has custom settings
        has_custom = settings_store._has_custom_settings(settings, 1)

        # Assert: Matches old default by base name, so not custom
        assert has_custom is False


@pytest.mark.asyncio
async def test_has_custom_settings_with_old_version_but_custom_model(settings_store):
    # Arrange: User with old version but custom model
    with (
        patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'),
        patch('server.constants.CURRENT_USER_SETTINGS_VERSION', 5),
        patch(
            'server.constants.USER_SETTINGS_VERSION_TO_MODEL',
            {1: 'claude-3-5-sonnet-20241022'},
        ),
    ):
        settings = Settings(llm_model='gpt-4', llm_base_url='http://default.url')

        # Act: Check if has custom settings
        has_custom = settings_store._has_custom_settings(settings, 1)

        # Assert: Custom model detected
        assert has_custom is True


@pytest.mark.asyncio
async def test_has_custom_settings_with_current_version(settings_store):
    # Arrange: User with current version
    with (
        patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'),
        patch('server.constants.CURRENT_USER_SETTINGS_VERSION', 5),
        patch(
            'server.constants.USER_SETTINGS_VERSION_TO_MODEL',
            {1: 'claude-3-5-sonnet-20241022', 5: 'claude-opus-4-5-20251101'},
        ),
    ):
        settings = Settings(
            llm_model='claude-3-5-sonnet-20241022', llm_base_url='http://default.url'
        )

        # Act: Check if has custom settings
        has_custom = settings_store._has_custom_settings(settings, 5)

        # Assert: Current version, so model is custom (not old default)
        assert has_custom is True


@pytest.mark.asyncio
async def test_has_custom_settings_with_none_version(settings_store):
    # Arrange: User with no version
    with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
        settings = Settings(
            llm_model='claude-3-5-sonnet-20241022', llm_base_url='http://default.url'
        )

        # Act: Check if has custom settings
        has_custom = settings_store._has_custom_settings(settings, None)

        # Assert: No version, so model is custom
        assert has_custom is True


@pytest.mark.asyncio
async def test_has_custom_settings_with_invalid_version(settings_store):
    # Arrange: User with invalid version
    with (
        patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'),
        patch('server.constants.CURRENT_USER_SETTINGS_VERSION', 5),
        patch(
            'server.constants.USER_SETTINGS_VERSION_TO_MODEL',
            {1: 'claude-3-5-sonnet-20241022'},
        ),
    ):
        settings = Settings(
            llm_model='claude-3-5-sonnet-20241022', llm_base_url='http://default.url'
        )

        # Act: Check if has custom settings
        has_custom = settings_store._has_custom_settings(settings, 99)

        # Assert: Invalid version, so model is custom
        assert has_custom is True


@pytest.mark.asyncio
async def test_has_custom_settings_normalizes_whitespace(settings_store):
    # Arrange: Settings with whitespace in values
    with patch('storage.saas_settings_store.LITE_LLM_API_URL', 'http://default.url'):
        settings = Settings(
            llm_model='  claude-3-5-sonnet-20241022  ',
            llm_base_url='  http://default.url  ',
        )

        # Act: Check if has custom settings
        has_custom = settings_store._has_custom_settings(settings, None)

        # Assert: Whitespace is normalized, custom model detected
        assert has_custom is True


@pytest.mark.asyncio
async def test_update_settings_upgrades_user_from_old_defaults(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User with old version using old defaults
    old_version = 1
    old_model = 'litellm_proxy/prod/claude-3-5-sonnet-20241022'
    settings = Settings(llm_model=old_model, llm_base_url=LITE_LLM_API_URL)

    # Use a consistent test URL
    test_base_url = 'http://test.url'

    with (
        patch('storage.saas_settings_store.session_maker', session_maker),
        patch(
            'server.constants.USER_SETTINGS_VERSION_TO_MODEL',
            {1: 'claude-3-5-sonnet-20241022', 5: 'claude-opus-4-5-20251101'},
        ),
        patch(
            'storage.saas_settings_store.USER_SETTINGS_VERSION_TO_MODEL',
            {1: 'claude-3-5-sonnet-20241022', 5: 'claude-opus-4-5-20251101'},
        ),
        patch('server.constants.CURRENT_USER_SETTINGS_VERSION', 5),
        patch('storage.saas_settings_store.CURRENT_USER_SETTINGS_VERSION', 5),
        patch('storage.saas_settings_store.LITE_LLM_API_URL', test_base_url),
        patch(
            'storage.saas_settings_store.get_default_litellm_model',
            return_value='litellm_proxy/prod/claude-opus-4-5-20251101',
        ),
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
    ):
        # Create existing user settings with old version
        with session_maker() as session:
            existing_settings = UserSettings(
                keycloak_user_id=settings_store.user_id,
                user_version=old_version,
                llm_model=old_model,
                llm_base_url=test_base_url,
            )
            session.add(existing_settings)
            session.commit()

        # Update settings to use test_base_url
        # Set user_version to match the database so _has_custom_settings can detect old defaults
        settings = Settings(
            llm_model=old_model, llm_base_url=test_base_url, user_version=old_version
        )

        # Act: Update settings
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

        # Assert: Settings upgraded to new defaults
        assert updated_settings is not None
        assert (
            updated_settings.llm_model == 'litellm_proxy/prod/claude-opus-4-5-20251101'
        )
        assert updated_settings.llm_base_url == test_base_url


@pytest.mark.asyncio
async def test_update_settings_preserves_custom_settings_during_upgrade(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User with old version but custom settings
    old_version = 1
    custom_model = 'gpt-4'
    custom_base_url = 'http://custom.url'
    settings = Settings(llm_model=custom_model, llm_base_url=custom_base_url)

    with (
        patch('storage.saas_settings_store.session_maker', session_maker),
        patch(
            'server.constants.USER_SETTINGS_VERSION_TO_MODEL',
            {1: 'claude-3-5-sonnet-20241022'},
        ),
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
    ):
        # Create existing user settings with old version
        with session_maker() as session:
            existing_settings = UserSettings(
                keycloak_user_id=settings_store.user_id,
                user_version=old_version,
                llm_model=custom_model,
                llm_base_url=custom_base_url,
            )
            session.add(existing_settings)
            session.commit()

        # Act: Update settings
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

        # Assert: Custom settings preserved
        assert updated_settings is not None
        assert updated_settings.llm_model == custom_model
        assert updated_settings.llm_base_url == custom_base_url


@pytest.mark.asyncio
async def test_update_settings_migrates_billing_margin_v3_to_v4(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User with version 3 and billing margin
    old_version = 3
    billing_margin = 2.0
    max_budget = 10.0
    spend = 5.0

    settings = Settings()

    mock_get_response = AsyncMock()
    mock_get_response.is_success = True
    mock_get_response.json = MagicMock(
        return_value={'user_info': {'max_budget': max_budget, 'spend': spend}}
    )

    mock_post_response = AsyncMock()
    mock_post_response.is_success = True
    mock_post_response.json = MagicMock(return_value={'key': 'test_api_key'})

    with (
        patch('storage.saas_settings_store.session_maker', session_maker),
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
        patch('httpx.AsyncClient') as mock_client,
    ):
        mock_client.return_value.__aenter__.return_value.get.return_value = (
            mock_get_response
        )
        mock_client.return_value.__aenter__.return_value.post.return_value = (
            mock_post_response
        )

        # Create existing user settings with version 3 and billing margin
        with session_maker() as session:
            existing_settings = UserSettings(
                keycloak_user_id=settings_store.user_id,
                user_version=old_version,
                billing_margin=billing_margin,
            )
            session.add(existing_settings)
            session.commit()

        # Act: Update settings
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

        # Assert: Settings updated
        assert updated_settings is not None

        # Assert: Billing margin applied to budget
        call_args = mock_client.return_value.__aenter__.return_value.post.call_args[1]
        assert call_args['json']['max_budget'] == max_budget * billing_margin
        assert call_args['json']['spend'] == spend * billing_margin

        # Assert: Billing margin reset to 1.0
        with session_maker() as session:
            updated_user_settings = (
                session.query(UserSettings)
                .filter(UserSettings.keycloak_user_id == settings_store.user_id)
                .first()
            )
            assert updated_user_settings.billing_margin == 1.0


@pytest.mark.asyncio
async def test_update_settings_skips_billing_margin_migration_when_already_v4(
    settings_store, mock_litellm_api, session_maker
):
    # Arrange: User with version 4
    version = 4
    billing_margin = 2.0
    max_budget = 10.0
    spend = 5.0

    settings = Settings()

    mock_get_response = AsyncMock()
    mock_get_response.is_success = True
    mock_get_response.json = MagicMock(
        return_value={'user_info': {'max_budget': max_budget, 'spend': spend}}
    )

    mock_post_response = AsyncMock()
    mock_post_response.is_success = True
    mock_post_response.json = MagicMock(return_value={'key': 'test_api_key'})

    with (
        patch('storage.saas_settings_store.session_maker', session_maker),
        patch(
            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
            AsyncMock(return_value={'email': 'user@example.com'}),
        ),
        patch('httpx.AsyncClient') as mock_client,
    ):
        mock_client.return_value.__aenter__.return_value.get.return_value = (
            mock_get_response
        )
        mock_client.return_value.__aenter__.return_value.post.return_value = (
            mock_post_response
        )

        # Create existing user settings with version 4
        with session_maker() as session:
            existing_settings = UserSettings(
                keycloak_user_id=settings_store.user_id,
                user_version=version,
                billing_margin=billing_margin,
            )
            session.add(existing_settings)
            session.commit()

        # Act: Update settings
        updated_settings = await settings_store.update_settings_with_litellm_default(
            settings
        )

        # Assert: Settings updated
        assert updated_settings is not None

        # Assert: Billing margin NOT applied (version >= 4)
        call_args = mock_client.return_value.__aenter__.return_value.post.call_args[1]
        assert call_args['json']['max_budget'] == max_budget
        assert call_args['json']['spend'] == spend