178 lines
6 KiB
Python
178 lines
6 KiB
Python
"""Tests for health tracker and router profile self-recovery from AUTH_REQUIRED state."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from antigravity_provider.router.adapters.antigravity_adapter import AntigravityAdapter
|
|
from antigravity_provider.router.health_tracker import (
|
|
AUTH_REQUIRED,
|
|
HEALTHY,
|
|
HealthTracker,
|
|
)
|
|
from antigravity_provider.router.profile_manager import ProfileAuthManager
|
|
from antigravity_provider.router.router_config import (
|
|
RolePolicy,
|
|
RouterConfig,
|
|
RouterProfileConfig,
|
|
)
|
|
from antigravity_provider.router.router_engine import RouterEngine
|
|
|
|
|
|
@pytest.fixture
|
|
def isolated_env(tmp_path, monkeypatch):
|
|
monkeypatch.setenv("HERMES_HOME", str(tmp_path))
|
|
state_file = tmp_path / "router_state.json"
|
|
ht = HealthTracker(state_file=state_file)
|
|
return tmp_path, ht
|
|
|
|
|
|
def test_auth_required_skipped_in_routing(isolated_env):
|
|
tmp_path, ht = isolated_env
|
|
|
|
# Configure 2 profiles for role 'manager'
|
|
cfg = RouterConfig(
|
|
profiles={
|
|
"ag-orch-1": RouterProfileConfig(
|
|
profile_id="ag-orch-1",
|
|
provider="antigravity",
|
|
enabled=True,
|
|
),
|
|
"ag-orch-2": RouterProfileConfig(
|
|
profile_id="ag-orch-2",
|
|
provider="antigravity",
|
|
enabled=True,
|
|
),
|
|
},
|
|
roles={
|
|
"manager": RolePolicy(
|
|
role_name="manager",
|
|
preferred_chain=["ag-orch-1", "ag-orch-2"],
|
|
),
|
|
},
|
|
default_role="manager",
|
|
)
|
|
|
|
engine = RouterEngine(config=cfg, health=ht)
|
|
|
|
def _mock_invoke(*args, **kwargs):
|
|
return {"id": "res-1", "choices": [{"message": {"role": "assistant", "content": "hello"}}]}
|
|
|
|
# Initial state: ag-orch-1 is healthy and selected
|
|
with patch.object(AntigravityAdapter, "invoke", side_effect=_mock_invoke):
|
|
res = engine.route_request({"prompt": "hi"}, role="manager")
|
|
assert res.get("router_metadata", {}).get("profile_id") == "ag-orch-1"
|
|
|
|
# Mark ag-orch-1 as AUTH_REQUIRED
|
|
ht.mark_auth_required("ag-orch-1", reason="401 Unauthorized token expired")
|
|
assert not ht.is_healthy("ag-orch-1")
|
|
|
|
# Routing should skip ag-orch-1 and route to ag-orch-2
|
|
with patch.object(AntigravityAdapter, "invoke", side_effect=_mock_invoke):
|
|
res = engine.route_request({"prompt": "hi"}, role="manager")
|
|
assert res.get("router_metadata", {}).get("profile_id") == "ag-orch-2"
|
|
|
|
|
|
def test_save_profile_auth_auto_recovery(isolated_env):
|
|
tmp_path, ht = isolated_env
|
|
|
|
cfg = RouterConfig(
|
|
profiles={
|
|
"ag-orch-1": RouterProfileConfig(
|
|
profile_id="ag-orch-1",
|
|
provider="antigravity",
|
|
enabled=True,
|
|
),
|
|
"ag-orch-2": RouterProfileConfig(
|
|
profile_id="ag-orch-2",
|
|
provider="antigravity",
|
|
enabled=True,
|
|
),
|
|
},
|
|
roles={
|
|
"manager": RolePolicy(
|
|
role_name="manager",
|
|
preferred_chain=["ag-orch-1", "ag-orch-2"],
|
|
),
|
|
},
|
|
default_role="manager",
|
|
)
|
|
engine = RouterEngine(config=cfg, health=ht)
|
|
|
|
def _mock_invoke(*args, **kwargs):
|
|
return {"id": "res-1", "choices": [{"message": {"role": "assistant", "content": "hello"}}]}
|
|
|
|
# Mark ag-orch-1 as AUTH_REQUIRED
|
|
ht.mark_auth_required("ag-orch-1", reason="Auth token rejected")
|
|
rec = ht.get_or_create("ag-orch-1")
|
|
assert rec.overall_state == AUTH_REQUIRED
|
|
assert rec.last_error == "Auth token rejected"
|
|
assert not ht.is_healthy("ag-orch-1")
|
|
|
|
# Saving credentials via ProfileAuthManager publishes event and auto-recovers health
|
|
ProfileAuthManager.save_profile_auth(
|
|
"antigravity",
|
|
"ag-orch-1",
|
|
{"token": {"access_token": "valid_tok_123", "refresh_token": "ref_123"}},
|
|
)
|
|
|
|
# State is automatically restored to HEALTHY without manual clear_cooldown
|
|
rec = ht.get_or_create("ag-orch-1")
|
|
assert rec.overall_state == HEALTHY
|
|
assert rec.last_error is None
|
|
assert ht.is_healthy("ag-orch-1")
|
|
|
|
# Router now routes back to primary profile ag-orch-1
|
|
with patch.object(AntigravityAdapter, "invoke", side_effect=_mock_invoke):
|
|
res = engine.route_request({"prompt": "hi"}, role="manager")
|
|
assert res.get("router_metadata", {}).get("profile_id") == "ag-orch-1"
|
|
|
|
|
|
def test_auth_file_mtime_recovery_in_is_healthy(isolated_env):
|
|
tmp_path, ht = isolated_env
|
|
|
|
# Setup profile directory and initial auth file
|
|
pdir = tmp_path / "agy_profiles" / "ag-test-1"
|
|
pdir.mkdir(parents=True, exist_ok=True)
|
|
auth_file = pdir / "auth.json"
|
|
auth_file.write_text(json.dumps({"token": "old"}), encoding="utf-8")
|
|
time.sleep(0.05)
|
|
|
|
# Mark AUTH_REQUIRED at a timestamp
|
|
ht.mark_auth_required("ag-test-1", reason="OAuth error")
|
|
rec = ht.get_or_create("ag-test-1")
|
|
assert rec.overall_state == AUTH_REQUIRED
|
|
assert not ht.is_healthy("ag-test-1")
|
|
|
|
# Now update auth.json with newer mtime
|
|
time.sleep(0.05)
|
|
auth_file.write_text(json.dumps({"token": {"access_token": "new_refreshed_token"}}), encoding="utf-8")
|
|
|
|
# Calling is_healthy should detect the newer valid auth file and auto-recover
|
|
assert ht.is_healthy("ag-test-1")
|
|
rec = ht.get_or_create("ag-test-1")
|
|
assert rec.overall_state == HEALTHY
|
|
assert rec.last_error is None
|
|
|
|
|
|
def test_clear_cooldown_full_reset(isolated_env):
|
|
tmp_path, ht = isolated_env
|
|
|
|
ht.mark_auth_required("ag-1", reason="Invalid grant")
|
|
rec = ht.get_or_create("ag-1")
|
|
assert rec.overall_state == AUTH_REQUIRED
|
|
assert rec.last_error == "Invalid grant"
|
|
|
|
# clear_cooldown should completely reset overall_state and last_error
|
|
ht.clear_cooldown("ag-1")
|
|
rec = ht.get_or_create("ag-1")
|
|
assert rec.overall_state == HEALTHY
|
|
assert rec.last_error is None
|
|
for frec in rec.families.values():
|
|
assert frec.state == HEALTHY
|
|
assert frec.last_error is None
|
|
assert frec.reason is None
|