learningcircuit--local-deep-research/tests/metrics/pricing/test_pricing_fetcher_coverage.py

"""
Comprehensive coverage tests for PricingFetcher.

Focuses on mocked HTTP interactions, edge cases, and all code paths:
- __init__, __aenter__, __aexit__ (with and without session)
- get_model_pricing: local providers, provider-first match, slash model
  name, exact match, no match, None inputs
- _get_models_by_provider: all provider branches
- get_provider_from_model: all prefix branches including ollama-only path
"""

import pytest
from unittest.mock import AsyncMock, MagicMock, patch

from local_deep_research.metrics.pricing.pricing_fetcher import PricingFetcher


# ===================================================================
# 1. Initialization
# ===================================================================
class TestInit:
    def test_session_starts_none(self):
        fetcher = PricingFetcher()
        assert fetcher.session is None

    def test_static_pricing_populated(self):
        fetcher = PricingFetcher()
        assert len(fetcher.static_pricing) > 0
        assert "gpt-4" in fetcher.static_pricing
        assert "ollama" in fetcher.static_pricing


# ===================================================================
# 2. Async context manager
# ===================================================================
class TestAsyncContextManager:
    @pytest.mark.asyncio
    async def test_aenter_creates_session(self):
        """__aenter__ should set self.session to an aiohttp.ClientSession."""
        with patch(
            "local_deep_research.metrics.pricing.pricing_fetcher.aiohttp.ClientSession"
        ) as mock_cls:
            mock_session = MagicMock()
            mock_cls.return_value = mock_session

            fetcher = PricingFetcher()
            result = await fetcher.__aenter__()

            assert result is fetcher
            assert fetcher.session is mock_session

    @pytest.mark.asyncio
    async def test_aexit_closes_session(self):
        """__aexit__ should close the session when one exists."""
        fetcher = PricingFetcher()
        mock_session = AsyncMock()
        fetcher.session = mock_session

        await fetcher.__aexit__(None, None, None)

        mock_session.close.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_aexit_no_session(self):
        """__aexit__ should be safe when session is None."""
        fetcher = PricingFetcher()
        assert fetcher.session is None
        # Should not raise
        await fetcher.__aexit__(None, None, None)


# ===================================================================
# 3. get_model_pricing -- local providers
# ===================================================================
class TestGetModelPricingLocalProviders:
    @pytest.mark.asyncio
    async def test_ollama_returns_zero(self):
        fetcher = PricingFetcher()
        result = await fetcher.get_model_pricing("anything", "ollama")
        assert result == {"prompt": 0.0, "completion": 0.0}

    @pytest.mark.asyncio
    async def test_lmstudio_returns_zero(self):
        fetcher = PricingFetcher()
        result = await fetcher.get_model_pricing("my-model", "lmstudio")
        assert result == {"prompt": 0.0, "completion": 0.0}

    @pytest.mark.asyncio
    async def test_llamacpp_returns_zero(self):
        fetcher = PricingFetcher()
        result = await fetcher.get_model_pricing("my-model", "llamacpp")
        assert result == {"prompt": 0.0, "completion": 0.0}


# ===================================================================
# 4. get_model_pricing -- provider-first match via _get_models_by_provider
# ===================================================================
class TestGetModelPricingProviderFirstMatch:
    @pytest.mark.asyncio
    async def test_provider_exact_match(self):
        """When provider is given, _get_models_by_provider is consulted."""
        fetcher = PricingFetcher()
        result = await fetcher.get_model_pricing("gpt-4", "openai")
        assert result is not None
        assert result["prompt"] == 0.03

    @pytest.mark.asyncio
    async def test_provider_slash_model_match(self):
        """Model names with slash strip the prefix in provider lookup."""
        fetcher = PricingFetcher()
        result = await fetcher.get_model_pricing("openai/gpt-4o", "openai")
        assert result is not None
        assert result["prompt"] == 0.005

    @pytest.mark.asyncio
    async def test_slash_model_without_provider(self):
        """Slash model name without provider falls to static_pricing."""
        fetcher = PricingFetcher()
        result = await fetcher.get_model_pricing("openai/gpt-4o-mini")
        assert result is not None
        assert result["prompt"] == 0.00015


# ===================================================================
# 5. get_model_pricing -- exact match and no match
# ===================================================================
class TestGetModelPricingExactAndNoMatch:
    @pytest.mark.asyncio
    async def test_exact_match_in_static_pricing(self):
        fetcher = PricingFetcher()
        result = await fetcher.get_model_pricing("claude-3-5-sonnet")
        assert result == {"prompt": 0.003, "completion": 0.015}

    @pytest.mark.asyncio
    async def test_no_match_returns_none(self):
        fetcher = PricingFetcher()
        result = await fetcher.get_model_pricing("totally-unknown-model-xyz")
        assert result is None

    @pytest.mark.asyncio
    async def test_none_model_name_returns_none(self):
        """None model name normalised to empty string returns None."""
        fetcher = PricingFetcher()
        result = await fetcher.get_model_pricing(None)
        assert result is None

    @pytest.mark.asyncio
    async def test_none_provider_handled(self):
        """None provider should be normalised to empty string."""
        fetcher = PricingFetcher()
        result = await fetcher.get_model_pricing("gpt-4", None)
        assert result is not None

    @pytest.mark.asyncio
    async def test_case_insensitive_model(self):
        """Model names should be lowercased for matching."""
        fetcher = PricingFetcher()
        result = await fetcher.get_model_pricing("GPT-4")
        assert result is not None
        assert result["prompt"] == 0.03


# ===================================================================
# 6. _get_models_by_provider
# ===================================================================
class TestGetModelsByProvider:
    def test_openai_returns_gpt_models(self):
        fetcher = PricingFetcher()
        models = fetcher._get_models_by_provider("openai")
        assert all(k.startswith("gpt") for k in models)
        assert "gpt-4" in models

    def test_anthropic_returns_claude_models(self):
        fetcher = PricingFetcher()
        models = fetcher._get_models_by_provider("anthropic")
        assert all(k.startswith("claude") for k in models)
        assert "claude-3-opus" in models

    def test_google_returns_gemini_models(self):
        fetcher = PricingFetcher()
        models = fetcher._get_models_by_provider("google")
        assert all(k.startswith("gemini") for k in models)
        assert "gemini-pro" in models

    def test_ollama_returns_free_models(self):
        fetcher = PricingFetcher()
        models = fetcher._get_models_by_provider("ollama")
        for pricing in models.values():
            assert pricing["prompt"] == 0.0
            assert pricing["completion"] == 0.0

    def test_unknown_provider_returns_empty(self):
        fetcher = PricingFetcher()
        models = fetcher._get_models_by_provider("some_random_provider")
        assert models == {}

    def test_case_insensitive_provider(self):
        fetcher = PricingFetcher()
        models = fetcher._get_models_by_provider("OpenAI")
        assert "gpt-4" in models


# ===================================================================
# 7. get_provider_from_model
# ===================================================================
class TestGetProviderFromModel:
    def test_gpt_returns_openai(self):
        fetcher = PricingFetcher()
        assert fetcher.get_provider_from_model("gpt-4o") == "openai"

    def test_openai_keyword_returns_openai(self):
        fetcher = PricingFetcher()
        assert fetcher.get_provider_from_model("openai-custom") == "openai"

    def test_claude_returns_anthropic(self):
        fetcher = PricingFetcher()
        assert fetcher.get_provider_from_model("claude-3-haiku") == "anthropic"

    def test_anthropic_keyword_returns_anthropic(self):
        fetcher = PricingFetcher()
        assert fetcher.get_provider_from_model("anthropic-v1") == "anthropic"

    def test_gemini_returns_google(self):
        fetcher = PricingFetcher()
        assert fetcher.get_provider_from_model("gemini-1.5-pro") == "google"

    def test_google_keyword_returns_google(self):
        fetcher = PricingFetcher()
        assert fetcher.get_provider_from_model("google-bard") == "google"

    def test_llama_returns_meta(self):
        fetcher = PricingFetcher()
        assert fetcher.get_provider_from_model("llama-3-8b") == "meta"

    def test_meta_keyword_returns_meta(self):
        fetcher = PricingFetcher()
        assert fetcher.get_provider_from_model("meta-llama-3") == "meta"

    def test_mistral_returns_mistral(self):
        fetcher = PricingFetcher()
        assert (
            fetcher.get_provider_from_model("mistral-7b-instruct") == "mistral"
        )

    def test_ollama_contains_llama_returns_meta(self):
        """'ollama' contains 'llama' so it matches meta first."""
        fetcher = PricingFetcher()
        assert fetcher.get_provider_from_model("ollama-server") == "meta"

    def test_unknown_model_returns_unknown(self):
        fetcher = PricingFetcher()
        assert fetcher.get_provider_from_model("phi-3-mini") == "unknown"

    def test_case_insensitive(self):
        fetcher = PricingFetcher()
        assert fetcher.get_provider_from_model("GPT-4-TURBO") == "openai"
        assert fetcher.get_provider_from_model("CLAUDE-3-OPUS") == "anthropic"