Torbjørn Lindahl
5add230d4c
Some checks failed
Semantic Release / release (push) Has been cancelled
Implements a new provider that uses Google's Antigravity unified gateway API to access Claude, Gemini, and other models through a single OAuth2-authenticated endpoint. Features: - OAuth2 token management with automatic refresh - Multi-account rotation for rate limit distribution - Support for Claude Opus/Sonnet 4.5 (with/without thinking) - Support for Gemini 2.5/3 models (Pro/Flash variants) - Thinking mode support with configurable tokens - Image processing support - Dual quota pool tracking (antigravity vs gemini-cli) - Gemini-style API request format Authentication: - Reads from ANTIGRAVITY_REFRESH_TOKEN env var (priority) - Falls back to ~/.config/opencode/antigravity-accounts.json - Automatic token refresh with retry logic - Rate limit tracking per account and quota pool Files added: - providers/antigravity.py - Main provider implementation - providers/antigravity_auth.py - OAuth token manager - providers/registries/antigravity.py - Model registry - conf/antigravity_models.json - Model definitions (11 models) - docs/antigravity_provider.md - Setup and usage docs - tests/test_antigravity_provider.py - Unit tests (14 pass) Integration: - Added to provider priority order after ZEN - Registered in server.py with auto-detection - ToS warning logged on first use
312 lines
12 KiB
Python
312 lines
12 KiB
Python
"""Unit tests for the Antigravity provider."""
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from providers.antigravity import AntigravityProvider
|
|
from providers.antigravity_auth import AntigravityAccount, AntigravityTokenManager
|
|
from providers.registries.antigravity import AntigravityModelRegistry
|
|
from providers.shared import ProviderType
|
|
|
|
|
|
class TestAntigravityModelRegistry:
|
|
"""Tests for the Antigravity model registry."""
|
|
|
|
def test_registry_loads_models(self):
|
|
"""Verify registry loads models from config file."""
|
|
registry = AntigravityModelRegistry()
|
|
models = registry.list_models()
|
|
|
|
assert len(models) > 0
|
|
assert "claude-opus-4-5-thinking" in models
|
|
assert "claude-sonnet-4-5" in models
|
|
assert "gemini-3-pro-high" in models
|
|
|
|
def test_registry_aliases(self):
|
|
"""Verify alias resolution works correctly."""
|
|
registry = AntigravityModelRegistry()
|
|
|
|
# Test alias resolution
|
|
capabilities = registry.resolve("antigravity-claude-opus-4-5-thinking")
|
|
assert capabilities is not None
|
|
assert capabilities.model_name == "claude-opus-4-5-thinking"
|
|
|
|
# Test short alias
|
|
capabilities = registry.resolve("ag-opus-thinking")
|
|
assert capabilities is not None
|
|
assert capabilities.model_name == "claude-opus-4-5-thinking"
|
|
|
|
def test_model_capabilities(self):
|
|
"""Verify model capabilities are correct."""
|
|
registry = AntigravityModelRegistry()
|
|
|
|
opus = registry.get_capabilities("claude-opus-4-5-thinking")
|
|
assert opus is not None
|
|
assert opus.supports_extended_thinking is True
|
|
assert opus.max_thinking_tokens == 32768
|
|
assert opus.supports_images is True
|
|
|
|
sonnet = registry.get_capabilities("claude-sonnet-4-5")
|
|
assert sonnet is not None
|
|
assert sonnet.supports_extended_thinking is False
|
|
|
|
def test_quota_pool(self):
|
|
"""Verify quota pool assignment."""
|
|
registry = AntigravityModelRegistry()
|
|
|
|
# Antigravity quota models
|
|
assert registry.get_quota_pool("claude-opus-4-5-thinking") == "antigravity"
|
|
assert registry.get_quota_pool("gemini-3-pro-high") == "antigravity"
|
|
|
|
# Gemini CLI quota models
|
|
assert registry.get_quota_pool("gemini-2.5-pro") == "gemini-cli"
|
|
assert registry.get_quota_pool("gemini-3-flash-preview") == "gemini-cli"
|
|
|
|
|
|
class TestAntigravityTokenManager:
|
|
"""Tests for the OAuth2 token manager."""
|
|
|
|
def test_empty_initialization(self):
|
|
"""Verify manager handles no accounts gracefully."""
|
|
# Create temp dir without accounts file
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
fake_path = Path(tmpdir) / "nonexistent.json"
|
|
manager = AntigravityTokenManager(accounts_file=fake_path)
|
|
assert not manager.has_accounts()
|
|
assert manager.get_account_count() == 0
|
|
|
|
def test_env_var_override(self, monkeypatch):
|
|
"""Verify environment variable takes precedence."""
|
|
monkeypatch.setenv("ANTIGRAVITY_REFRESH_TOKEN", "test_token_123")
|
|
monkeypatch.setenv("ANTIGRAVITY_PROJECT_ID", "test-project")
|
|
|
|
manager = AntigravityTokenManager()
|
|
assert manager.has_accounts()
|
|
assert manager.get_account_count() == 1
|
|
|
|
def test_load_accounts_from_file(self):
|
|
"""Verify loading accounts from JSON file."""
|
|
accounts_data = {
|
|
"version": 3,
|
|
"accounts": [
|
|
{
|
|
"email": "test@example.com",
|
|
"refreshToken": "token_1",
|
|
"projectId": "project-1",
|
|
"enabled": True,
|
|
},
|
|
{
|
|
"email": "test2@example.com",
|
|
"refreshToken": "token_2",
|
|
"projectId": "project-2",
|
|
"enabled": True,
|
|
},
|
|
{
|
|
"email": "disabled@example.com",
|
|
"refreshToken": "token_3",
|
|
"projectId": "project-3",
|
|
"enabled": False, # Should be skipped
|
|
},
|
|
],
|
|
}
|
|
|
|
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
|
|
json.dump(accounts_data, f)
|
|
temp_path = f.name
|
|
|
|
try:
|
|
# Clear env to ensure file is used
|
|
with mock.patch.dict(os.environ, {}, clear=True):
|
|
manager = AntigravityTokenManager(accounts_file=temp_path)
|
|
assert manager.has_accounts()
|
|
assert manager.get_account_count() == 2 # Disabled account skipped
|
|
finally:
|
|
os.unlink(temp_path)
|
|
|
|
|
|
class TestAntigravityAccount:
|
|
"""Tests for the AntigravityAccount dataclass."""
|
|
|
|
def test_token_validity(self):
|
|
"""Test token validity checking."""
|
|
import time
|
|
|
|
account = AntigravityAccount(
|
|
email="test@example.com",
|
|
refresh_token="token",
|
|
project_id="project",
|
|
)
|
|
|
|
# No access token
|
|
assert not account.is_token_valid()
|
|
|
|
# Set valid token
|
|
account.access_token = "access_token"
|
|
account.token_expiry = time.time() + 3600
|
|
assert account.is_token_valid()
|
|
|
|
# Expired token
|
|
account.token_expiry = time.time() - 100
|
|
assert not account.is_token_valid()
|
|
|
|
def test_rate_limit_tracking(self):
|
|
"""Test rate limit tracking."""
|
|
import time
|
|
|
|
account = AntigravityAccount(
|
|
email="test@example.com",
|
|
refresh_token="token",
|
|
project_id="project",
|
|
)
|
|
|
|
# No rate limits initially
|
|
assert not account.is_rate_limited()
|
|
assert not account.is_rate_limited("claude-opus-4-5-thinking")
|
|
|
|
# Set rate limit
|
|
account.set_rate_limited("claude-opus-4-5-thinking", 60.0)
|
|
assert account.is_rate_limited("claude-opus-4-5-thinking")
|
|
assert account.is_rate_limited() # Any rate limit active
|
|
|
|
|
|
class TestAntigravityProvider:
|
|
"""Tests for the Antigravity provider."""
|
|
|
|
def test_provider_type(self):
|
|
"""Verify provider type is correct."""
|
|
# Mock token manager to avoid needing real credentials
|
|
with mock.patch.object(AntigravityTokenManager, "__init__", return_value=None):
|
|
with mock.patch.object(AntigravityTokenManager, "has_accounts", return_value=True):
|
|
provider = AntigravityProvider()
|
|
assert provider.get_provider_type() == ProviderType.ANTIGRAVITY
|
|
|
|
def test_model_capabilities(self):
|
|
"""Verify model capabilities are accessible."""
|
|
with mock.patch.object(AntigravityTokenManager, "__init__", return_value=None):
|
|
with mock.patch.object(AntigravityTokenManager, "has_accounts", return_value=True):
|
|
provider = AntigravityProvider()
|
|
caps = provider.get_all_model_capabilities()
|
|
assert len(caps) > 0
|
|
assert "claude-opus-4-5-thinking" in caps
|
|
|
|
def test_model_resolution(self):
|
|
"""Verify alias resolution works through provider."""
|
|
with mock.patch.object(AntigravityTokenManager, "__init__", return_value=None):
|
|
with mock.patch.object(AntigravityTokenManager, "has_accounts", return_value=True):
|
|
provider = AntigravityProvider()
|
|
|
|
# Test alias resolution
|
|
resolved = provider._resolve_model_name("ag-opus-thinking")
|
|
assert resolved == "claude-opus-4-5-thinking"
|
|
|
|
# Test canonical name
|
|
resolved = provider._resolve_model_name("claude-sonnet-4-5")
|
|
assert resolved == "claude-sonnet-4-5"
|
|
|
|
def test_request_building(self):
|
|
"""Verify request body is built correctly."""
|
|
with mock.patch.object(AntigravityTokenManager, "__init__", return_value=None):
|
|
with mock.patch.object(AntigravityTokenManager, "has_accounts", return_value=True):
|
|
provider = AntigravityProvider()
|
|
|
|
capabilities = provider.get_capabilities("claude-opus-4-5-thinking")
|
|
|
|
request = provider._build_request(
|
|
prompt="Hello, world!",
|
|
model_name="claude-opus-4-5-thinking",
|
|
project_id="test-project",
|
|
system_prompt="You are a helpful assistant.",
|
|
temperature=0.7,
|
|
max_output_tokens=1000,
|
|
thinking_mode="medium",
|
|
capabilities=capabilities,
|
|
)
|
|
|
|
assert request["project"] == "test-project"
|
|
assert request["model"] == "claude-opus-4-5-thinking"
|
|
assert "contents" in request["request"]
|
|
assert "systemInstruction" in request["request"]
|
|
assert "generationConfig" in request["request"]
|
|
|
|
# Check thinking config for thinking models
|
|
config = request["request"]["generationConfig"]
|
|
assert "thinkingConfig" in config
|
|
assert config["thinkingConfig"]["thinkingBudget"] > 0
|
|
|
|
def test_list_models(self):
|
|
"""Verify list_models returns expected models."""
|
|
with mock.patch.object(AntigravityTokenManager, "__init__", return_value=None):
|
|
with mock.patch.object(AntigravityTokenManager, "has_accounts", return_value=True):
|
|
provider = AntigravityProvider()
|
|
models = provider.list_models(respect_restrictions=False)
|
|
|
|
assert len(models) > 0
|
|
# Check that both canonical names and aliases are present
|
|
assert any("claude" in m.lower() for m in models)
|
|
assert any("gemini" in m.lower() for m in models)
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestAntigravityIntegration:
|
|
"""Integration tests (require actual credentials and valid tokens).
|
|
|
|
These tests require:
|
|
- Valid Antigravity credentials (env var or accounts file)
|
|
- Working OAuth tokens that can be refreshed
|
|
|
|
Run with: pytest -m integration tests/test_antigravity_provider.py
|
|
Skip with: pytest -m "not integration" tests/test_antigravity_provider.py
|
|
"""
|
|
|
|
@pytest.mark.skipif(
|
|
not os.getenv("ANTIGRAVITY_REFRESH_TOKEN")
|
|
and not Path.home().joinpath(".config/opencode/antigravity-accounts.json").exists(),
|
|
reason="Antigravity credentials not available",
|
|
)
|
|
def test_token_refresh(self):
|
|
"""Test actual token refresh with real credentials."""
|
|
manager = AntigravityTokenManager()
|
|
assert manager.has_accounts()
|
|
|
|
# This will trigger a token refresh - may fail if tokens are expired
|
|
try:
|
|
token, project_id, headers = manager.get_access_token()
|
|
assert token is not None
|
|
assert len(token) > 0
|
|
assert project_id is not None
|
|
except RuntimeError as e:
|
|
if "rate limited or have invalid tokens" in str(e):
|
|
pytest.skip("Antigravity tokens are expired or invalid")
|
|
raise
|
|
|
|
@pytest.mark.skipif(
|
|
not os.getenv("ANTIGRAVITY_REFRESH_TOKEN")
|
|
and not Path.home().joinpath(".config/opencode/antigravity-accounts.json").exists(),
|
|
reason="Antigravity credentials not available",
|
|
)
|
|
def test_simple_generation(self):
|
|
"""Test actual content generation with Antigravity API."""
|
|
provider = AntigravityProvider()
|
|
|
|
try:
|
|
response = provider.generate_content(
|
|
prompt="Say 'hello' and nothing else.",
|
|
model_name="claude-sonnet-4-5",
|
|
temperature=0.0,
|
|
max_output_tokens=100,
|
|
)
|
|
except RuntimeError as e:
|
|
if "rate limited or have invalid tokens" in str(e):
|
|
pytest.skip("Antigravity tokens are expired or invalid")
|
|
raise
|
|
|
|
assert response is not None
|
|
assert response.content is not None
|
|
assert "hello" in response.content.lower()
|
|
assert response.provider == ProviderType.ANTIGRAVITY
|