wehub-resource-sync
121 lines
4.0 KiB
Python
121 lines
4.0 KiB
Python
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from mlflow.assistant.config import AssistantConfig, PermissionsConfig
|
|
from mlflow.assistant.providers import MlflowGatewayProvider, OllamaProvider
|
|
from mlflow.assistant.providers.base import clear_config_cache
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def config_file(tmp_path):
|
|
config_path = tmp_path / "config.json"
|
|
clear_config_cache()
|
|
with patch("mlflow.assistant.config.CONFIG_PATH", config_path):
|
|
yield config_path
|
|
clear_config_cache()
|
|
|
|
|
|
def test_update_provider_preserves_selection():
|
|
ollama = OllamaProvider.OLLAMA_PROVIDER_NAME
|
|
config = AssistantConfig()
|
|
config.set_provider("claude_code", "claude-4", base_url=None)
|
|
assert config.providers["claude_code"].selected is True
|
|
|
|
config.update_provider(ollama, model="llama3.2", base_url="http://localhost:11434")
|
|
|
|
assert config.providers["claude_code"].selected is True
|
|
assert config.providers[ollama].selected is False
|
|
assert config.providers[ollama].model == "llama3.2"
|
|
assert config.providers[ollama].base_url == "http://localhost:11434"
|
|
|
|
|
|
def test_update_provider_creates_entry_when_missing():
|
|
ollama = OllamaProvider.OLLAMA_PROVIDER_NAME
|
|
config = AssistantConfig()
|
|
assert ollama not in config.providers
|
|
|
|
config.update_provider(ollama, model="llama3.2", base_url="http://localhost:11434")
|
|
|
|
assert ollama in config.providers
|
|
assert config.providers[ollama].model == "llama3.2"
|
|
assert config.providers[ollama].selected is False
|
|
|
|
|
|
def test_update_provider_does_not_deselect_existing():
|
|
ollama = OllamaProvider.OLLAMA_PROVIDER_NAME
|
|
config = AssistantConfig()
|
|
config.set_provider(ollama, "llama3.2")
|
|
assert config.providers[ollama].selected is True
|
|
|
|
config.update_provider(ollama, base_url="http://newhost:11434")
|
|
|
|
assert config.providers[ollama].selected is True
|
|
assert config.providers[ollama].base_url == "http://newhost:11434"
|
|
|
|
|
|
def test_set_provider_deselects_others():
|
|
ollama = OllamaProvider.OLLAMA_PROVIDER_NAME
|
|
config = AssistantConfig()
|
|
config.set_provider("claude_code", "claude-4")
|
|
config.update_provider(ollama, model="llama3.2")
|
|
|
|
config.set_provider(ollama, "llama3.2")
|
|
|
|
assert config.providers["claude_code"].selected is False
|
|
assert config.providers[ollama].selected is True
|
|
|
|
|
|
def test_get_selected_provider_returns_none_when_empty():
|
|
config = AssistantConfig()
|
|
assert config.get_selected_provider() is None
|
|
|
|
|
|
def test_config_round_trip(config_file):
|
|
ollama = OllamaProvider.OLLAMA_PROVIDER_NAME
|
|
config = AssistantConfig()
|
|
config.set_provider(ollama, "llama3.2", base_url="http://localhost:11434")
|
|
config.save()
|
|
|
|
loaded = AssistantConfig.load()
|
|
assert loaded.providers[ollama].model == "llama3.2"
|
|
assert loaded.providers[ollama].selected is True
|
|
assert loaded.providers[ollama].base_url == "http://localhost:11434"
|
|
|
|
|
|
def test_update_provider_preserves_permissions():
|
|
ollama = OllamaProvider.OLLAMA_PROVIDER_NAME
|
|
config = AssistantConfig()
|
|
config.set_provider(ollama, "llama3.2")
|
|
new_perms = PermissionsConfig(allow_edit_files=False, full_access=False)
|
|
|
|
config.update_provider(ollama, permissions=new_perms)
|
|
|
|
assert config.providers[ollama].permissions.allow_edit_files is False
|
|
|
|
|
|
def test_api_key_round_trip(config_file):
|
|
config = AssistantConfig()
|
|
config.set_provider(
|
|
MlflowGatewayProvider.GATEWAY_PROVIDER_NAME,
|
|
"gpt-4",
|
|
base_url="http://gateway.local:5000",
|
|
api_key="sk-test-123",
|
|
)
|
|
config.save()
|
|
|
|
loaded = AssistantConfig.load()
|
|
name = MlflowGatewayProvider.GATEWAY_PROVIDER_NAME
|
|
assert loaded.providers[name].api_key == "sk-test-123"
|
|
assert loaded.providers[name].base_url == "http://gateway.local:5000"
|
|
|
|
|
|
def test_update_provider_updates_api_key():
|
|
name = MlflowGatewayProvider.GATEWAY_PROVIDER_NAME
|
|
config = AssistantConfig()
|
|
config.set_provider(name, "gpt-4", api_key="sk-old")
|
|
|
|
config.update_provider(name, api_key="sk-new")
|
|
|
|
assert config.providers[name].api_key == "sk-new"
|