hermes-agent/tests/agent/test_gemini_free_tier_gate.py

"""Tests for Gemini free-tier detection and blocking."""
from __future__ import annotations

from unittest.mock import MagicMock, patch


from agent.gemini_native_adapter import (
    gemini_http_error,
    is_free_tier_quota_error,
    probe_gemini_tier,
)


def _mock_response(status: int, headers: dict | None = None, text: str = "") -> MagicMock:
    resp = MagicMock()
    resp.status_code = status
    resp.headers = headers or {}
    resp.text = text
    return resp


def _run_probe(resp: MagicMock) -> str:
    with patch("agent.gemini_native_adapter.httpx.Client") as MC:
        inst = MagicMock()
        inst.post.return_value = resp
        MC.return_value.__enter__.return_value = inst
        return probe_gemini_tier("fake-key")


class TestProbeGeminiTier:
    """Verify the tier probe classifies keys correctly."""

    def test_free_tier_via_rpd_header_flash(self):
        # gemini-2.5-flash free tier: 250 RPD
        resp = _mock_response(200, {"x-ratelimit-limit-requests-per-day": "250"}, "{}")
        assert _run_probe(resp) == "free"

    def test_free_tier_via_rpd_header_pro(self):
        # gemini-2.5-pro free tier: 100 RPD
        resp = _mock_response(200, {"x-ratelimit-limit-requests-per-day": "100"}, "{}")
        assert _run_probe(resp) == "free"

    def test_free_tier_via_rpd_header_flash_lite(self):
        # flash-lite free tier: 1000 RPD (our upper bound)
        resp = _mock_response(200, {"x-ratelimit-limit-requests-per-day": "1000"}, "{}")
        assert _run_probe(resp) == "free"

    def test_paid_tier_via_rpd_header(self):
        # Tier 1 starts at 1500+ RPD
        resp = _mock_response(200, {"x-ratelimit-limit-requests-per-day": "1500"}, "{}")
        assert _run_probe(resp) == "paid"

    def test_free_tier_via_429_body(self):
        body = (
            '{"error":{"code":429,"message":"Quota exceeded for metric: '
            'generativelanguage.googleapis.com/generate_content_free_tier_requests, '
            'limit: 20"}}'
        )
        resp = _mock_response(429, {}, body)
        assert _run_probe(resp) == "free"

    def test_paid_429_has_no_free_tier_marker(self):
        body = '{"error":{"code":429,"message":"rate limited"}}'
        resp = _mock_response(429, {}, body)
        assert _run_probe(resp) == "paid"

    def test_successful_200_without_rpd_header_is_paid(self):
        resp = _mock_response(200, {}, '{"candidates":[]}')
        assert _run_probe(resp) == "paid"

    def test_401_returns_unknown(self):
        resp = _mock_response(401, {}, '{"error":{"code":401}}')
        assert _run_probe(resp) == "unknown"

    def test_404_returns_unknown(self):
        resp = _mock_response(404, {}, '{"error":{"code":404}}')
        assert _run_probe(resp) == "unknown"

    def test_network_error_returns_unknown(self):
        with patch(
            "agent.gemini_native_adapter.httpx.Client",
            side_effect=Exception("dns failure"),
        ):
            assert probe_gemini_tier("fake-key") == "unknown"

    def test_empty_key_returns_unknown(self):
        assert probe_gemini_tier("") == "unknown"
        assert probe_gemini_tier("   ") == "unknown"
        assert probe_gemini_tier(None) == "unknown"  # type: ignore[arg-type]

    def test_malformed_rpd_header_falls_through(self):
        # Non-integer header value shouldn't crash; 200 with no usable header -> paid.
        resp = _mock_response(200, {"x-ratelimit-limit-requests-per-day": "abc"}, "{}")
        assert _run_probe(resp) == "paid"

    def test_openai_compat_suffix_stripped(self):
        """Base URLs ending in /openai get normalized to the native endpoint."""
        resp = _mock_response(200, {"x-ratelimit-limit-requests-per-day": "1500"}, "{}")
        with patch("agent.gemini_native_adapter.httpx.Client") as MC:
            inst = MagicMock()
            inst.post.return_value = resp
            MC.return_value.__enter__.return_value = inst
            probe_gemini_tier(
                "fake",
                "https://generativelanguage.googleapis.com/v1beta/openai",
            )
            # Verify the post URL does NOT contain /openai
            called_url = inst.post.call_args[0][0]
            assert "/openai/" not in called_url
            assert called_url.endswith(":generateContent")


class TestIsFreeTierQuotaError:
    def test_detects_free_tier_marker(self):
        assert is_free_tier_quota_error(
            "Quota exceeded for metric: generate_content_free_tier_requests"
        )

    def test_case_insensitive(self):
        assert is_free_tier_quota_error("QUOTA: FREE_TIER_REQUESTS")

    def test_no_free_tier_marker(self):
        assert not is_free_tier_quota_error("rate limited")

    def test_empty_string(self):
        assert not is_free_tier_quota_error("")

    def test_none(self):
        assert not is_free_tier_quota_error(None)  # type: ignore[arg-type]


class TestGeminiHttpErrorFreeTierGuidance:
    """gemini_http_error should append free-tier guidance for free-tier 429s."""

    class _FakeResp:
        def __init__(self, status: int, text: str):
            self.status_code = status
            self.headers: dict = {}
            self.text = text

    def test_free_tier_429_appends_guidance(self):
        body = (
            '{"error":{"code":429,"message":"Quota exceeded for metric: '
            "generativelanguage.googleapis.com/generate_content_free_tier_requests, "
            'limit: 20","status":"RESOURCE_EXHAUSTED"}}'
        )
        err = gemini_http_error(self._FakeResp(429, body))
        msg = str(err)
        assert "free tier" in msg.lower()
        assert "aistudio.google.com/apikey" in msg

    def test_paid_429_has_no_billing_url(self):
        body = '{"error":{"code":429,"message":"Rate limited","status":"RESOURCE_EXHAUSTED"}}'
        err = gemini_http_error(self._FakeResp(429, body))
        assert "aistudio.google.com/apikey" not in str(err)

    def test_non_429_has_no_billing_url(self):
        body = '{"error":{"code":400,"message":"bad request","status":"INVALID_ARGUMENT"}}'
        err = gemini_http_error(self._FakeResp(400, body))
        assert "aistudio.google.com/apikey" not in str(err)

    def test_401_has_no_billing_url(self):
        body = '{"error":{"code":401,"message":"API key invalid","status":"UNAUTHENTICATED"}}'
        err = gemini_http_error(self._FakeResp(401, body))
        assert "aistudio.google.com/apikey" not in str(err)