242 lines
8.7 KiB
Python
242 lines
8.7 KiB
Python
import os
|
|
from dataclasses import dataclass
|
|
from typing import Optional
|
|
|
|
import pytest
|
|
|
|
from ra_aid.env import validate_environment
|
|
|
|
|
|
@dataclass
|
|
class MockArgs:
|
|
provider: str
|
|
expert_provider: str
|
|
model: Optional[str] = None
|
|
expert_model: Optional[str] = None
|
|
research_provider: Optional[str] = None
|
|
research_model: Optional[str] = None
|
|
planner_provider: Optional[str] = None
|
|
planner_model: Optional[str] = None
|
|
|
|
|
|
@pytest.fixture
|
|
def clean_env(monkeypatch):
|
|
"""Remove relevant environment variables before each test"""
|
|
env_vars = [
|
|
"ANTHROPIC_API_KEY",
|
|
"OPENAI_API_KEY",
|
|
"OPENROUTER_API_KEY",
|
|
"OPENAI_API_BASE",
|
|
"EXPERT_ANTHROPIC_API_KEY",
|
|
"EXPERT_OPENAI_API_KEY",
|
|
"EXPERT_OPENROUTER_API_KEY",
|
|
"EXPERT_OPENAI_API_BASE",
|
|
"TAVILY_API_KEY",
|
|
"ANTHROPIC_MODEL",
|
|
]
|
|
for var in env_vars:
|
|
monkeypatch.delenv(var, raising=False)
|
|
|
|
|
|
def test_anthropic_validation(clean_env, monkeypatch):
|
|
args = MockArgs(
|
|
provider="anthropic", expert_provider="openai", model="claude-3-haiku-20240307"
|
|
)
|
|
|
|
# Should fail without API key
|
|
with pytest.raises(SystemExit):
|
|
validate_environment(args)
|
|
|
|
# Should pass with API key and model
|
|
monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key")
|
|
expert_enabled, expert_missing, web_research_enabled, web_research_missing = (
|
|
validate_environment(args)
|
|
)
|
|
assert not expert_enabled
|
|
assert not web_research_enabled
|
|
assert "TAVILY_API_KEY environment variable is not set" in web_research_missing
|
|
|
|
|
|
def test_openai_validation(clean_env, monkeypatch):
|
|
args = MockArgs(provider="openai", expert_provider="openai")
|
|
|
|
# Should fail without API key
|
|
with pytest.raises(SystemExit):
|
|
validate_environment(args)
|
|
|
|
# Should pass with API key and enable expert mode with fallback
|
|
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
|
|
expert_enabled, expert_missing, web_research_enabled, web_research_missing = (
|
|
validate_environment(args)
|
|
)
|
|
assert expert_enabled
|
|
assert not expert_missing
|
|
assert not web_research_enabled
|
|
assert "TAVILY_API_KEY environment variable is not set" in web_research_missing
|
|
assert os.environ.get("EXPERT_OPENAI_API_KEY") == "test-key"
|
|
|
|
|
|
def test_openai_compatible_validation(clean_env, monkeypatch):
|
|
args = MockArgs(provider="openai-compatible", expert_provider="openai-compatible")
|
|
|
|
# Should fail without API key and base URL
|
|
with pytest.raises(SystemExit):
|
|
validate_environment(args)
|
|
|
|
# Should fail with only API key
|
|
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
|
|
with pytest.raises(SystemExit):
|
|
validate_environment(args)
|
|
|
|
# Should pass with both API key and base URL
|
|
monkeypatch.setenv("OPENAI_API_BASE", "http://test")
|
|
expert_enabled, expert_missing, web_research_enabled, web_research_missing = (
|
|
validate_environment(args)
|
|
)
|
|
assert expert_enabled
|
|
assert not expert_missing
|
|
assert not web_research_enabled
|
|
assert "TAVILY_API_KEY environment variable is not set" in web_research_missing
|
|
assert os.environ.get("EXPERT_OPENAI_API_KEY") == "test-key"
|
|
assert os.environ.get("EXPERT_OPENAI_API_BASE") == "http://test"
|
|
|
|
|
|
def test_expert_fallback(clean_env, monkeypatch):
|
|
args = MockArgs(provider="openai", expert_provider="openai")
|
|
|
|
# Set only base API key
|
|
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
|
|
|
|
# Should enable expert mode with fallback
|
|
expert_enabled, expert_missing, web_research_enabled, web_research_missing = (
|
|
validate_environment(args)
|
|
)
|
|
assert expert_enabled
|
|
assert not expert_missing
|
|
assert not web_research_enabled
|
|
assert "TAVILY_API_KEY environment variable is not set" in web_research_missing
|
|
assert os.environ.get("EXPERT_OPENAI_API_KEY") == "test-key"
|
|
|
|
# Should use explicit expert key if available
|
|
monkeypatch.setenv("EXPERT_OPENAI_API_KEY", "expert-key")
|
|
expert_enabled, expert_missing, web_research_enabled, web_research_missing = (
|
|
validate_environment(args)
|
|
)
|
|
assert expert_enabled
|
|
assert not expert_missing
|
|
assert not web_research_enabled
|
|
assert "TAVILY_API_KEY environment variable is not set" in web_research_missing
|
|
assert os.environ.get("EXPERT_OPENAI_API_KEY") == "expert-key"
|
|
|
|
|
|
def test_cross_provider_fallback(clean_env, monkeypatch):
|
|
"""Test that fallback works even when providers differ"""
|
|
args = MockArgs(
|
|
provider="openai",
|
|
expert_provider="anthropic",
|
|
expert_model="claude-3-haiku-20240307",
|
|
)
|
|
|
|
# Set base API key for main provider and expert provider
|
|
monkeypatch.setenv("OPENAI_API_KEY", "openai-key")
|
|
monkeypatch.setenv("ANTHROPIC_API_KEY", "anthropic-key")
|
|
monkeypatch.setenv("ANTHROPIC_MODEL", "claude-3-haiku-20240307")
|
|
|
|
# Should enable expert mode with fallback to ANTHROPIC base key
|
|
expert_enabled, expert_missing, web_research_enabled, web_research_missing = (
|
|
validate_environment(args)
|
|
)
|
|
assert expert_enabled
|
|
assert not expert_missing
|
|
assert not web_research_enabled
|
|
assert "TAVILY_API_KEY environment variable is not set" in web_research_missing
|
|
|
|
# Try with openai-compatible expert provider
|
|
args = MockArgs(provider="anthropic", expert_provider="openai-compatible")
|
|
monkeypatch.setenv("OPENAI_API_KEY", "openai-key")
|
|
monkeypatch.setenv("OPENAI_API_BASE", "http://test")
|
|
|
|
expert_enabled, expert_missing, web_research_enabled, web_research_missing = (
|
|
validate_environment(args)
|
|
)
|
|
assert expert_enabled
|
|
assert not expert_missing
|
|
assert not web_research_enabled
|
|
assert "TAVILY_API_KEY environment variable is not set" in web_research_missing
|
|
assert os.environ.get("EXPERT_OPENAI_API_KEY") == "openai-key"
|
|
assert os.environ.get("EXPERT_OPENAI_API_BASE") == "http://test"
|
|
|
|
|
|
def test_no_warning_on_fallback(clean_env, monkeypatch):
|
|
"""Test that no warning is issued when fallback succeeds"""
|
|
args = MockArgs(provider="openai", expert_provider="openai")
|
|
|
|
# Set only base API key
|
|
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
|
|
|
|
# Should enable expert mode with fallback and no warnings
|
|
expert_enabled, expert_missing, web_research_enabled, web_research_missing = (
|
|
validate_environment(args)
|
|
)
|
|
assert expert_enabled
|
|
assert not expert_missing
|
|
assert not web_research_enabled
|
|
assert "TAVILY_API_KEY environment variable is not set" in web_research_missing
|
|
assert os.environ.get("EXPERT_OPENAI_API_KEY") == "test-key"
|
|
|
|
# Should use explicit expert key if available
|
|
monkeypatch.setenv("EXPERT_OPENAI_API_KEY", "expert-key")
|
|
expert_enabled, expert_missing, web_research_enabled, web_research_missing = (
|
|
validate_environment(args)
|
|
)
|
|
assert expert_enabled
|
|
assert not expert_missing
|
|
assert not web_research_enabled
|
|
assert "TAVILY_API_KEY environment variable is not set" in web_research_missing
|
|
assert os.environ.get("EXPERT_OPENAI_API_KEY") == "expert-key"
|
|
|
|
|
|
def test_different_providers_no_expert_key(clean_env, monkeypatch):
|
|
"""Test behavior when providers differ and only base keys are available"""
|
|
args = MockArgs(
|
|
provider="anthropic", expert_provider="openai", model="claude-3-haiku-20240307"
|
|
)
|
|
|
|
# Set only base keys
|
|
monkeypatch.setenv("ANTHROPIC_API_KEY", "anthropic-key")
|
|
monkeypatch.setenv("OPENAI_API_KEY", "openai-key")
|
|
|
|
# Should enable expert mode and use base OPENAI key
|
|
expert_enabled, expert_missing, web_research_enabled, web_research_missing = (
|
|
validate_environment(args)
|
|
)
|
|
assert expert_enabled
|
|
assert not expert_missing
|
|
assert not web_research_enabled
|
|
assert "TAVILY_API_KEY environment variable is not set" in web_research_missing
|
|
|
|
|
|
def test_mixed_provider_openai_compatible(clean_env, monkeypatch):
|
|
"""Test behavior with openai-compatible expert and different main provider"""
|
|
args = MockArgs(
|
|
provider="anthropic",
|
|
expert_provider="openai-compatible",
|
|
model="claude-3-haiku-20240307",
|
|
)
|
|
|
|
# Set all required keys and URLs
|
|
monkeypatch.setenv("ANTHROPIC_API_KEY", "anthropic-key")
|
|
monkeypatch.setenv("OPENAI_API_KEY", "openai-key")
|
|
monkeypatch.setenv("OPENAI_API_BASE", "http://test")
|
|
|
|
# Should enable expert mode and use base openai key and URL
|
|
expert_enabled, expert_missing, web_research_enabled, web_research_missing = (
|
|
validate_environment(args)
|
|
)
|
|
assert expert_enabled
|
|
assert not expert_missing
|
|
assert not web_research_enabled
|
|
assert "TAVILY_API_KEY environment variable is not set" in web_research_missing
|
|
assert os.environ.get("EXPERT_OPENAI_API_KEY") == "openai-key"
|
|
assert os.environ.get("EXPERT_OPENAI_API_BASE") == "http://test"
|