tests/test_ccr_context_tracker.py

"""Tests for CCR context tracker.

These tests verify that:
1. Compression events are tracked correctly
2. Query analysis finds relevant compressed contexts
3. Proactive expansion works as expected
4. LRU eviction and TTL work correctly
5. Expansion recommendations are appropriate
"""

import json
import time

import pytest

from headroom.cache.compression_store import (
    get_compression_store,
    reset_compression_store,
)
from headroom.ccr.context_tracker import (
    CompressedContext,
    ContextTracker,
    ContextTrackerConfig,
    ExpansionRecommendation,
    get_context_tracker,
    reset_context_tracker,
)


class TestContextTrackerBasics:
    """Test basic context tracking functionality."""

    @pytest.fixture(autouse=True)
    def reset_trackers(self):
        """Reset trackers before each test."""
        reset_context_tracker()
        reset_compression_store()
        yield
        reset_context_tracker()
        reset_compression_store()

    def test_track_compression(self):
        """Track a compression event."""
        tracker = ContextTracker()

        tracker.track_compression(
            hash_key="abc123",
            turn_number=1,
            tool_name="Bash",
            original_count=100,
            compressed_count=10,
            query_context="find all python files",
            sample_content='["src/main.py", "src/auth.py"]',
        )

        assert "abc123" in tracker.get_tracked_hashes()
        stats = tracker.get_stats()
        assert stats["tracked_contexts"] == 1

    def test_track_multiple_compressions(self):
        """Track multiple compression events."""
        tracker = ContextTracker()

        for i in range(5):
            tracker.track_compression(
                hash_key=f"hash_{i}",
                turn_number=i,
                tool_name="Bash",
                original_count=100,
                compressed_count=10,
            )

        hashes = tracker.get_tracked_hashes()
        assert len(hashes) == 5
        assert all(f"hash_{i}" in hashes for i in range(5))

    def test_tracking_disabled(self):
        """Tracking disabled doesn't store contexts."""
        config = ContextTrackerConfig(enabled=False)
        tracker = ContextTracker(config)

        tracker.track_compression(
            hash_key="abc123",
            turn_number=1,
            tool_name="Bash",
            original_count=100,
            compressed_count=10,
        )

        assert len(tracker.get_tracked_hashes()) == 0

    def test_update_existing_hash(self):
        """Updating existing hash updates the context data."""
        tracker = ContextTracker()

        tracker.track_compression(
            hash_key="first", turn_number=1, tool_name=None, original_count=50, compressed_count=5
        )
        tracker.track_compression(
            hash_key="second", turn_number=2, tool_name=None, original_count=50, compressed_count=5
        )
        tracker.track_compression(
            hash_key="first", turn_number=3, tool_name=None, original_count=60, compressed_count=6
        )

        # Should still have 2 unique hashes
        hashes = tracker.get_tracked_hashes()
        assert len(hashes) == 2
        assert "first" in hashes
        assert "second" in hashes

        # The updated context should have the new values
        stats = tracker.get_stats()
        first_ctx = next(c for c in stats["contexts"] if c["hash"] == "first")
        assert first_ctx["items"] == "6/60"  # Updated values
        assert first_ctx["turn"] == 3


class TestLRUEviction:
    """Test LRU eviction at capacity."""

    def test_eviction_at_capacity(self):
        """Oldest entries evicted when at capacity."""
        config = ContextTrackerConfig(max_tracked_contexts=3)
        tracker = ContextTracker(config)

        for i in range(5):
            tracker.track_compression(
                hash_key=f"hash_{i}",
                turn_number=i,
                tool_name=None,
                original_count=100,
                compressed_count=10,
            )
            time.sleep(0.01)  # Ensure different timestamps

        hashes = tracker.get_tracked_hashes()
        assert len(hashes) == 3
        # Should have the last 3
        assert "hash_2" in hashes
        assert "hash_3" in hashes
        assert "hash_4" in hashes
        # First 2 should be evicted
        assert "hash_0" not in hashes
        assert "hash_1" not in hashes


class TestQueryAnalysis:
    """Test query analysis for relevance detection."""

    @pytest.fixture(autouse=True)
    def reset_trackers(self):
        """Reset trackers before each test."""
        reset_context_tracker()
        reset_compression_store()
        yield
        reset_context_tracker()
        reset_compression_store()

    def test_analyze_query_finds_relevant_context(self):
        """Query analysis finds relevant compressed context."""
        # Use lower relevance threshold to make test more reliable
        config = ContextTrackerConfig(relevance_threshold=0.1)
        tracker = ContextTracker(config)

        tracker.track_compression(
            hash_key="auth_hash",
            turn_number=1,
            tool_name="Bash",
            original_count=100,
            compressed_count=10,
            query_context="find authentication files",
            # Use more explicit content with keywords that will match
            sample_content="authentication middleware handler login security",
        )

        recommendations = tracker.analyze_query(
            query="show authentication middleware",
            current_turn=2,
        )

        assert len(recommendations) >= 1
        assert recommendations[0].hash_key == "auth_hash"

    def test_analyze_query_no_match(self):
        """Query analysis returns empty for unrelated query."""
        tracker = ContextTracker()

        tracker.track_compression(
            hash_key="db_hash",
            turn_number=1,
            tool_name="Bash",
            original_count=100,
            compressed_count=10,
            query_context="find database files",
            sample_content='["database.py", "models.py"]',
        )

        recommendations = tracker.analyze_query(
            query="What is the weather like?",
            current_turn=2,
        )

        # Should not match unrelated query
        assert len(recommendations) == 0

    def test_analyze_query_keyword_overlap(self):
        """Query matches based on keyword overlap."""
        tracker = ContextTracker()

        tracker.track_compression(
            hash_key="python_files",
            turn_number=1,
            tool_name="Glob",
            original_count=200,
            compressed_count=20,
            query_context="find python files",
            sample_content='["main.py", "utils.py", "config.py", "test_main.py"]',
        )

        # Query with overlapping keywords
        recommendations = tracker.analyze_query(
            query="Show me the main python file",
            current_turn=2,
        )

        assert len(recommendations) >= 1

    def test_analyze_query_proactive_disabled(self):
        """No recommendations when proactive expansion disabled."""
        config = ContextTrackerConfig(proactive_expansion=False)
        tracker = ContextTracker(config)

        tracker.track_compression(
            hash_key="abc123",
            turn_number=1,
            tool_name="Bash",
            original_count=100,
            compressed_count=10,
            sample_content='["relevant.py"]',
        )

        recommendations = tracker.analyze_query(
            query="Show me relevant files",
            current_turn=2,
        )

        assert len(recommendations) == 0

    def test_analyze_query_respects_age(self):
        """Old contexts get lower relevance scores."""
        config = ContextTrackerConfig(max_context_age_seconds=2.0)
        tracker = ContextTracker(config)

        tracker.track_compression(
            hash_key="old_context",
            turn_number=1,
            tool_name="Bash",
            original_count=100,
            compressed_count=10,
            sample_content='["auth.py"]',
        )

        # Wait for context to age
        time.sleep(2.1)

        recommendations = tracker.analyze_query(
            query="Show me the authentication code",
            current_turn=5,
        )

        # Should not recommend aged-out context
        assert len(recommendations) == 0

    def test_analyze_query_max_recommendations(self):
        """Respects max proactive expansions limit."""
        config = ContextTrackerConfig(max_proactive_expansions=2)
        tracker = ContextTracker(config)

        # Track many relevant contexts
        for i in range(5):
            tracker.track_compression(
                hash_key=f"hash_{i}",
                turn_number=i,
                tool_name="Bash",
                original_count=100,
                compressed_count=10,
                sample_content=f'["python_{i}.py", "main.py"]',
            )

        recommendations = tracker.analyze_query(
            query="Show me the python main file",
            current_turn=10,
        )

        assert len(recommendations) <= 2


class TestRelevanceCalculation:
    """Test relevance score calculation."""

    def test_extract_keywords(self):
        """Keywords are extracted correctly."""
        tracker = ContextTracker()

        keywords = tracker._extract_keywords("Find authentication middleware files")

        assert "authentication" in keywords
        assert "middleware" in keywords
        assert "files" in keywords
        # Stop words should be filtered
        assert "the" not in keywords

    def test_exact_substring_match_bonus(self):
        """Exact substring matches get bonus score."""
        tracker = ContextTracker()

        tracker.track_compression(
            hash_key="exact_match",
            turn_number=1,
            tool_name=None,
            original_count=100,
            compressed_count=10,
            sample_content="authentication_middleware.py, auth_handler.py",
        )

        # Query with exact substring match
        recommendations = tracker.analyze_query(
            query="authentication middleware",
            current_turn=2,
        )

        assert len(recommendations) >= 1
        # Should have high relevance
        assert recommendations[0].relevance_score > 0.3


class TestExpansionTypeDetection:
    """Test determination of expansion type (full vs search)."""

    def test_full_expansion_high_relevance(self):
        """High relevance triggers full expansion."""
        tracker = ContextTracker()

        context = CompressedContext(
            hash_key="test",
            turn_number=1,
            timestamp=time.time(),
            tool_name="Bash",
            original_item_count=50,
            compressed_item_count=5,
            query_context="find files",
            sample_content="auth.py, middleware.py",
        )

        expand_full, search_query = tracker._determine_expansion_type(
            query="authentication middleware",
            context=context,
            relevance=0.8,  # High relevance
        )

        assert expand_full is True
        assert search_query is None

    def test_full_expansion_small_count(self):
        """Small original item count triggers full expansion."""
        tracker = ContextTracker()

        context = CompressedContext(
            hash_key="test",
            turn_number=1,
            timestamp=time.time(),
            tool_name="Bash",
            original_item_count=30,  # Small
            compressed_item_count=5,
            query_context="find files",
            sample_content="file.py",
        )

        expand_full, search_query = tracker._determine_expansion_type(
            query="some query",
            context=context,
            relevance=0.4,
        )

        assert expand_full is True

    def test_search_expansion_large_count(self):
        """Large original count with specific keywords triggers search."""
        tracker = ContextTracker()

        context = CompressedContext(
            hash_key="test",
            turn_number=1,
            timestamp=time.time(),
            tool_name="Bash",
            original_item_count=500,  # Large
            compressed_item_count=20,
            query_context="find all files",
            sample_content="many files...",
        )

        expand_full, search_query = tracker._determine_expansion_type(
            query="find authentication middleware handler",
            context=context,
            relevance=0.4,  # Medium relevance
        )

        # Should use search for large datasets
        if not expand_full:
            assert search_query is not None
            assert "authentication" in search_query or "middleware" in search_query


class TestExpansionExecution:
    """Test execution of expansion recommendations."""

    @pytest.fixture(autouse=True)
    def reset_stores(self):
        """Reset stores before each test."""
        reset_context_tracker()
        reset_compression_store()
        yield
        reset_context_tracker()
        reset_compression_store()

    def test_execute_full_expansion(self):
        """Execute full expansion retrieval."""
        store = get_compression_store()
        original = json.dumps([{"id": i} for i in range(100)])

        hash_key = store.store(
            original=original,
            compressed="[]",
            original_item_count=100,
        )

        tracker = ContextTracker()
        recommendations = [
            ExpansionRecommendation(
                hash_key=hash_key,
                reason="relevant to query",
                relevance_score=0.8,
                expand_full=True,
            )
        ]

        results = tracker.execute_expansions(recommendations)

        assert len(results) == 1
        assert results[0]["type"] == "full"
        assert results[0]["item_count"] == 100

    def test_execute_search_expansion(self):
        """Execute search expansion."""
        store = get_compression_store()
        items = [
            {"id": 1, "content": "authentication code"},
            {"id": 2, "content": "database operations"},
            {"id": 3, "content": "authentication middleware"},
        ]
        original = json.dumps(items)

        hash_key = store.store(
            original=original,
            compressed="[]",
            original_item_count=3,
        )

        tracker = ContextTracker()
        recommendations = [
            ExpansionRecommendation(
                hash_key=hash_key,
                reason="relevant to query",
                relevance_score=0.5,
                expand_full=False,
                search_query="authentication",
            )
        ]

        results = tracker.execute_expansions(recommendations)

        assert len(results) == 1
        assert results[0]["type"] == "search"
        assert results[0]["query"] == "authentication"

    def test_execute_nonexistent_hash(self):
        """Handle expansion of nonexistent hash gracefully."""
        tracker = ContextTracker()
        recommendations = [
            ExpansionRecommendation(
                hash_key="nonexistent123",
                reason="test",
                relevance_score=0.5,
                expand_full=True,
            )
        ]

        results = tracker.execute_expansions(recommendations)

        # Should handle gracefully, no results
        assert len(results) == 0


class TestExpansionFormatting:
    """Test formatting of expansions for context."""

    def test_format_full_expansion(self):
        """Format full expansion for LLM context."""
        tracker = ContextTracker()

        expansions = [
            {
                "hash": "abc123",
                "type": "full",
                "content": '[{"id": 1}, {"id": 2}]',
                "item_count": 2,
                "reason": "relevant to query",
            }
        ]

        formatted = tracker.format_expansions_for_context(expansions)

        assert "[Proactive Context Expansion" in formatted
        assert "Expanded from earlier" in formatted
        assert '[{"id": 1}, {"id": 2}]' in formatted

    def test_format_search_expansion(self):
        """Format search expansion for LLM context."""
        tracker = ContextTracker()

        expansions = [
            {
                "hash": "def456",
                "type": "search",
                "query": "authentication",
                "content": [{"id": 1, "content": "auth"}],
                "item_count": 1,
                "reason": "matched query",
            }
        ]

        formatted = tracker.format_expansions_for_context(expansions)

        assert "Search results for 'authentication'" in formatted

    def test_format_empty_expansions(self):
        """Empty expansions return empty string."""
        tracker = ContextTracker()

        formatted = tracker.format_expansions_for_context([])

        assert formatted == ""


class TestGlobalTracker:
    """Test global tracker singleton."""

    @pytest.fixture(autouse=True)
    def reset_tracker(self):
        """Reset global tracker."""
        reset_context_tracker()
        yield
        reset_context_tracker()

    def test_singleton_pattern(self):
        """Global tracker uses singleton pattern."""
        tracker1 = get_context_tracker()
        tracker2 = get_context_tracker()

        assert tracker1 is tracker2

    def test_reset_clears_tracker(self):
        """Reset clears the global tracker."""
        tracker = get_context_tracker()
        tracker.track_compression(
            hash_key="test",
            turn_number=1,
            tool_name=None,
            original_count=10,
            compressed_count=1,
        )

        assert len(tracker.get_tracked_hashes()) == 1

        reset_context_tracker()
        new_tracker = get_context_tracker()

        assert len(new_tracker.get_tracked_hashes()) == 0


class TestContextTrackerConfig:
    """Test context tracker configuration."""

    def test_default_config(self):
        """Default config values."""
        config = ContextTrackerConfig()

        assert config.enabled is True
        assert config.max_tracked_contexts == 100
        assert config.relevance_threshold == 0.3
        assert config.max_context_age_seconds == 300.0
        assert config.proactive_expansion is True
        assert config.max_proactive_expansions == 2

    def test_custom_config(self):
        """Custom config values."""
        config = ContextTrackerConfig(
            enabled=False,
            max_tracked_contexts=50,
            relevance_threshold=0.5,
            max_proactive_expansions=5,
        )

        assert config.enabled is False
        assert config.max_tracked_contexts == 50
        assert config.relevance_threshold == 0.5
        assert config.max_proactive_expansions == 5


class TestCompressedContextDataClass:
    """Test CompressedContext dataclass."""

    def test_create_context(self):
        """Create compressed context."""
        context = CompressedContext(
            hash_key="abc123",
            turn_number=5,
            timestamp=1234567890.0,
            tool_name="Bash",
            original_item_count=100,
            compressed_item_count=10,
            query_context="find files",
            sample_content='["file1.py", "file2.py"]',
        )

        assert context.hash_key == "abc123"
        assert context.turn_number == 5
        assert context.tool_name == "Bash"
        assert context.original_item_count == 100
        assert context.compressed_item_count == 10


class TestExpansionRecommendationDataClass:
    """Test ExpansionRecommendation dataclass."""

    def test_full_expansion_recommendation(self):
        """Create full expansion recommendation."""
        rec = ExpansionRecommendation(
            hash_key="abc123",
            reason="high relevance",
            relevance_score=0.9,
            expand_full=True,
        )

        assert rec.expand_full is True
        assert rec.search_query is None

    def test_search_expansion_recommendation(self):
        """Create search expansion recommendation."""
        rec = ExpansionRecommendation(
            hash_key="def456",
            reason="partial match",
            relevance_score=0.5,
            expand_full=False,
            search_query="authentication",
        )

        assert rec.expand_full is False
        assert rec.search_query == "authentication"


class TestContextTrackerStats:
    """Test tracker statistics."""

    def test_stats_structure(self):
        """Stats have expected structure."""
        tracker = ContextTracker()

        tracker.track_compression(
            hash_key="test",
            turn_number=1,
            tool_name="Bash",
            original_count=100,
            compressed_count=10,
        )

        stats = tracker.get_stats()

        assert "tracked_contexts" in stats
        assert "current_turn" in stats
        assert "config" in stats
        assert "contexts" in stats

        assert stats["tracked_contexts"] == 1
        assert len(stats["contexts"]) == 1

    def test_stats_context_details(self):
        """Stats include context details."""
        tracker = ContextTracker()

        tracker.track_compression(
            hash_key="abc123",
            turn_number=3,
            tool_name="Glob",
            original_count=50,
            compressed_count=5,
        )

        stats = tracker.get_stats()
        context_stat = stats["contexts"][0]

        assert context_stat["hash"] == "abc123"
        assert context_stat["turn"] == 3
        assert context_stat["tool"] == "Glob"
        assert context_stat["items"] == "5/50"


class TestTrackerClear:
    """Test tracker clear functionality."""

    def test_clear_removes_all(self):
        """Clear removes all tracked contexts."""
        tracker = ContextTracker()

        for i in range(5):
            tracker.track_compression(
                hash_key=f"hash_{i}",
                turn_number=i,
                tool_name=None,
                original_count=10,
                compressed_count=1,
            )

        assert len(tracker.get_tracked_hashes()) == 5

        tracker.clear()

        assert len(tracker.get_tracked_hashes()) == 0
        stats = tracker.get_stats()
        assert stats["current_turn"] == 0