OpenHands/tests/unit/test_settings_store_functions.py
Rohit Malhotra 6d1e1f75ae
[Feat]: Support self hosted gitlab + enterprise github (#8274)
Co-authored-by: openhands <openhands@all-hands.dev>
Co-authored-by: amanape <83104063+amanape@users.noreply.github.com>
2025-05-08 17:54:32 -04:00

286 lines
9.1 KiB
Python

from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi.testclient import TestClient
from pydantic import SecretStr
from openhands.integrations.provider import ProviderToken
from openhands.integrations.service_types import ProviderType
from openhands.server.routes.secrets import (
app,
check_provider_tokens,
)
from openhands.server.routes.settings import store_llm_settings
from openhands.server.settings import POSTProviderModel
from openhands.storage import get_file_store
from openhands.storage.data_models.settings import Settings
from openhands.storage.data_models.user_secrets import UserSecrets
from openhands.storage.secrets.file_secrets_store import FileSecretsStore
# Mock functions to simulate the actual functions in settings.py
async def get_settings_store(request):
"""Mock function to get settings store."""
return MagicMock()
@pytest.fixture
def test_client():
# Create a test client
with (
patch(
'openhands.server.routes.secrets.check_provider_tokens',
AsyncMock(return_value=''),
),
):
client = TestClient(app)
yield client
@pytest.fixture
def temp_dir(tmp_path_factory: pytest.TempPathFactory) -> str:
return str(tmp_path_factory.mktemp('secrets_store'))
@pytest.fixture
def file_secrets_store(temp_dir):
file_store = get_file_store('local', temp_dir)
store = FileSecretsStore(file_store)
with patch(
'openhands.storage.secrets.file_secrets_store.FileSecretsStore.get_instance',
AsyncMock(return_value=store),
):
yield store
# Tests for check_provider_tokens
@pytest.mark.asyncio
async def test_check_provider_tokens_valid():
"""Test check_provider_tokens with valid tokens."""
provider_token = ProviderToken(token=SecretStr('valid-token'))
providers = POSTProviderModel(provider_tokens={ProviderType.GITHUB: provider_token})
# Empty existing provider tokens
existing_provider_tokens = {}
# Mock the validate_provider_token function to return GITHUB for valid tokens
with patch(
'openhands.server.routes.secrets.validate_provider_token'
) as mock_validate:
mock_validate.return_value = ProviderType.GITHUB
result = await check_provider_tokens(providers, existing_provider_tokens)
# Should return empty string for valid token
assert result == ''
mock_validate.assert_called_once()
@pytest.mark.asyncio
async def test_check_provider_tokens_invalid():
"""Test check_provider_tokens with invalid tokens."""
provider_token = ProviderToken(token=SecretStr('invalid-token'))
providers = POSTProviderModel(provider_tokens={ProviderType.GITHUB: provider_token})
# Empty existing provider tokens
existing_provider_tokens = {}
# Mock the validate_provider_token function to return None for invalid tokens
with patch(
'openhands.server.routes.secrets.validate_provider_token'
) as mock_validate:
mock_validate.return_value = None
result = await check_provider_tokens(providers, existing_provider_tokens)
# Should return error message for invalid token
assert 'Invalid token' in result
mock_validate.assert_called_once()
@pytest.mark.asyncio
async def test_check_provider_tokens_wrong_type():
"""Test check_provider_tokens with unsupported provider type."""
# We can't test with an unsupported provider type directly since the model enforces valid types
# Instead, we'll test with an empty provider_tokens dictionary
providers = POSTProviderModel(provider_tokens={})
# Empty existing provider tokens
existing_provider_tokens = {}
result = await check_provider_tokens(providers, existing_provider_tokens)
# Should return empty string for no providers
assert result == ''
@pytest.mark.asyncio
async def test_check_provider_tokens_no_tokens():
"""Test check_provider_tokens with no tokens."""
providers = POSTProviderModel(provider_tokens={})
# Empty existing provider tokens
existing_provider_tokens = {}
result = await check_provider_tokens(providers, existing_provider_tokens)
# Should return empty string when no tokens provided
assert result == ''
# Tests for store_llm_settings
@pytest.mark.asyncio
async def test_store_llm_settings_new_settings():
"""Test store_llm_settings with new settings."""
settings = Settings(
llm_model='gpt-4',
llm_api_key='test-api-key',
llm_base_url='https://api.example.com',
)
# Mock the settings store
mock_store = MagicMock()
mock_store.load = AsyncMock(return_value=None) # No existing settings
result = await store_llm_settings(settings, mock_store)
# Should return settings with the provided values
assert result.llm_model == 'gpt-4'
assert result.llm_api_key.get_secret_value() == 'test-api-key'
assert result.llm_base_url == 'https://api.example.com'
@pytest.mark.asyncio
async def test_store_llm_settings_update_existing():
"""Test store_llm_settings updates existing settings."""
settings = Settings(
llm_model='gpt-4',
llm_api_key='new-api-key',
llm_base_url='https://new.example.com',
)
# Mock the settings store
mock_store = MagicMock()
# Create existing settings
existing_settings = Settings(
llm_model='gpt-3.5',
llm_api_key=SecretStr('old-api-key'),
llm_base_url='https://old.example.com',
)
mock_store.load = AsyncMock(return_value=existing_settings)
result = await store_llm_settings(settings, mock_store)
# Should return settings with the updated values
assert result.llm_model == 'gpt-4'
assert result.llm_api_key.get_secret_value() == 'new-api-key'
assert result.llm_base_url == 'https://new.example.com'
@pytest.mark.asyncio
async def test_store_llm_settings_partial_update():
"""Test store_llm_settings with partial update."""
settings = Settings(
llm_model='gpt-4' # Only updating model
)
# Mock the settings store
mock_store = MagicMock()
# Create existing settings
existing_settings = Settings(
llm_model='gpt-3.5',
llm_api_key=SecretStr('existing-api-key'),
llm_base_url='https://existing.example.com',
)
mock_store.load = AsyncMock(return_value=existing_settings)
result = await store_llm_settings(settings, mock_store)
# Should return settings with updated model but keep other values
assert result.llm_model == 'gpt-4'
# For SecretStr objects, we need to compare the secret value
assert result.llm_api_key.get_secret_value() == 'existing-api-key'
assert result.llm_base_url == 'https://existing.example.com'
# Tests for store_provider_tokens
@pytest.mark.asyncio
async def test_store_provider_tokens_new_tokens(test_client, file_secrets_store):
"""Test store_provider_tokens with new tokens."""
provider_tokens = {'provider_tokens': {'github': {'token': 'new-token'}}}
# Mock the settings store
mock_store = MagicMock()
mock_store.load = AsyncMock(return_value=None) # No existing settings
UserSecrets()
user_secrets = await file_secrets_store.store(UserSecrets())
response = test_client.post('/api/add-git-providers', json=provider_tokens)
assert response.status_code == 200
user_secrets = await file_secrets_store.load()
assert (
user_secrets.provider_tokens[ProviderType.GITHUB].token.get_secret_value()
== 'new-token'
)
@pytest.mark.asyncio
async def test_store_provider_tokens_update_existing(test_client, file_secrets_store):
"""Test store_provider_tokens updates existing tokens."""
# Create existing settings with a GitHub token
github_token = ProviderToken(token=SecretStr('old-token'))
provider_tokens = {ProviderType.GITHUB: github_token}
# Create a UserSecrets with the provider tokens
user_secrets = UserSecrets(provider_tokens=provider_tokens)
await file_secrets_store.store(user_secrets)
response = test_client.post(
'/api/add-git-providers',
json={'provider_tokens': {'github': {'token': 'updated-token'}}},
)
assert response.status_code == 200
user_secrets = await file_secrets_store.load()
assert (
user_secrets.provider_tokens[ProviderType.GITHUB].token.get_secret_value()
== 'updated-token'
)
@pytest.mark.asyncio
async def test_store_provider_tokens_keep_existing(test_client, file_secrets_store):
"""Test store_provider_tokens keeps existing tokens when empty string provided."""
# Create existing secrets with a GitHub token
github_token = ProviderToken(token=SecretStr('existing-token'))
provider_tokens = {ProviderType.GITHUB: github_token}
user_secrets = UserSecrets(provider_tokens=provider_tokens)
await file_secrets_store.store(user_secrets)
response = test_client.post(
'/api/add-git-providers',
json={'provider_tokens': {'github': {'token': ''}}},
)
assert response.status_code == 200
user_secrets = await file_secrets_store.load()
assert (
user_secrets.provider_tokens[ProviderType.GITHUB].token.get_secret_value()
== 'existing-token'
)