headroom/tests/test_cache/test_client_integration.py

"""Tests for HeadroomClient cache optimizer integration."""

import os
import tempfile
from unittest.mock import MagicMock

import pytest

from headroom import (
    AnthropicCacheOptimizer,
    HeadroomClient,
)


@pytest.fixture
def temp_db():
    """Create a temporary database file."""
    fd, path = tempfile.mkstemp(suffix=".db")
    os.close(fd)
    yield f"sqlite:///{path}"
    if os.path.exists(path):
        os.unlink(path)


class MockTokenCounter:
    """Mock token counter for testing."""

    def count_tokens(self, text: str) -> int:
        return len(text) // 4

    def count_messages(self, messages: list) -> int:
        total = 0
        for msg in messages:
            content = msg.get("content", "")
            if isinstance(content, str):
                total += len(content) // 4
            elif isinstance(content, list):
                for block in content:
                    if isinstance(block, dict):
                        total += len(block.get("text", "")) // 4
        return total


class MockAnthropicProvider:
    """Mock Anthropic provider for testing."""

    name = "anthropic"

    def get_token_counter(self, model: str):
        return MockTokenCounter()

    def get_context_limit(self, model: str) -> int:
        return 200000


class MockOpenAIProvider:
    """Mock OpenAI provider for testing."""

    name = "openai"

    def get_token_counter(self, model: str):
        return MockTokenCounter()

    def get_context_limit(self, model: str) -> int:
        return 128000


class TestHeadroomClientCacheIntegration:
    """Test HeadroomClient cache optimizer integration."""

    def test_auto_detect_anthropic_optimizer(self, temp_db):
        """Test that Anthropic optimizer is auto-detected."""
        mock_client = MagicMock()
        provider = MockAnthropicProvider()

        client = HeadroomClient(
            original_client=mock_client,
            provider=provider,
            store_url=temp_db,
            enable_cache_optimizer=True,
        )

        assert client._cache_optimizer is not None
        assert client._cache_optimizer.name == "anthropic-cache-optimizer"

    def test_auto_detect_openai_optimizer(self, temp_db):
        """Test that OpenAI optimizer is auto-detected."""
        mock_client = MagicMock()
        provider = MockOpenAIProvider()

        client = HeadroomClient(
            original_client=mock_client,
            provider=provider,
            store_url=temp_db,
            enable_cache_optimizer=True,
        )

        assert client._cache_optimizer is not None
        assert client._cache_optimizer.name == "openai-prefix-stabilizer"

    def test_custom_optimizer(self, temp_db):
        """Test using a custom optimizer."""
        mock_client = MagicMock()
        provider = MockAnthropicProvider()
        custom_optimizer = AnthropicCacheOptimizer()

        client = HeadroomClient(
            original_client=mock_client,
            provider=provider,
            store_url=temp_db,
            cache_optimizer=custom_optimizer,
        )

        assert client._cache_optimizer is custom_optimizer

    def test_disable_cache_optimizer(self, temp_db):
        """Test disabling cache optimizer."""
        mock_client = MagicMock()
        provider = MockAnthropicProvider()

        client = HeadroomClient(
            original_client=mock_client,
            provider=provider,
            store_url=temp_db,
            enable_cache_optimizer=False,
        )

        assert client._cache_optimizer is None

    def test_semantic_cache_layer_creation(self, temp_db):
        """Test semantic cache layer is created when enabled."""
        mock_client = MagicMock()
        provider = MockAnthropicProvider()

        client = HeadroomClient(
            original_client=mock_client,
            provider=provider,
            store_url=temp_db,
            enable_cache_optimizer=True,
            enable_semantic_cache=True,
        )

        assert client._semantic_cache_layer is not None
        assert client._cache_optimizer is not None

    def test_extract_query_from_string_content(self, temp_db):
        """Test query extraction from string content."""
        mock_client = MagicMock()
        provider = MockAnthropicProvider()

        client = HeadroomClient(
            original_client=mock_client,
            provider=provider,
            store_url=temp_db,
        )

        messages = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "What is 2+2?"},
        ]

        query = client._extract_query(messages)
        assert query == "What is 2+2?"

    def test_extract_query_from_content_blocks(self, temp_db):
        """Test query extraction from content block format."""
        mock_client = MagicMock()
        provider = MockAnthropicProvider()

        client = HeadroomClient(
            original_client=mock_client,
            provider=provider,
            store_url=temp_db,
        )

        messages = [
            {"role": "system", "content": "You are helpful."},
            {
                "role": "user",
                "content": [{"type": "text", "text": "What is 2+2?"}],
            },
        ]

        query = client._extract_query(messages)
        assert query == "What is 2+2?"

    def test_extract_query_last_user_message(self, temp_db):
        """Test that query extraction uses last user message."""
        mock_client = MagicMock()
        provider = MockAnthropicProvider()

        client = HeadroomClient(
            original_client=mock_client,
            provider=provider,
            store_url=temp_db,
        )

        messages = [
            {"role": "user", "content": "First question"},
            {"role": "assistant", "content": "First answer"},
            {"role": "user", "content": "Second question"},
        ]

        query = client._extract_query(messages)
        assert query == "Second question"

    def test_config_propagation(self, temp_db):
        """Test that config is properly propagated."""
        mock_client = MagicMock()
        provider = MockAnthropicProvider()

        client = HeadroomClient(
            original_client=mock_client,
            provider=provider,
            store_url=temp_db,
            enable_cache_optimizer=True,
            enable_semantic_cache=True,
        )

        assert client._config.cache_optimizer.enabled is True
        assert client._config.cache_optimizer.enable_semantic_cache is True