tatlock/tests/core/test_benchmarks.py

"""
Tests for benchmark storage.

Tests performance tracking, Redis storage, and analytics features.
"""
import json
from datetime import datetime, timedelta, timezone
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from src.core.benchmarks import (
    BenchmarkStore,
    PerformanceBenchmark,
    get_benchmark_store,
)


class TestPerformanceBenchmark:
    """Test PerformanceBenchmark model."""

    def test_benchmark_creation(self):
        """Test creating a performance benchmark."""
        benchmark = PerformanceBenchmark(
            operation="steward_analysis",
            duration_seconds=1.23,
            success=True,
            recommendation_count=3,
        )

        assert benchmark.operation == "steward_analysis"
        assert benchmark.duration_seconds == 1.23
        assert benchmark.success is True
        assert benchmark.recommendation_count == 3
        assert isinstance(benchmark.timestamp, datetime)

    def test_benchmark_with_tool_fields(self):
        """Test benchmark with tool-specific fields."""
        benchmark = PerformanceBenchmark(
            operation="tool_call",
            duration_seconds=0.5,
            success=True,
            tool_name="calculate",
            was_recommended=True,
            was_actually_used=True,
        )

        assert benchmark.tool_name == "calculate"
        assert benchmark.was_recommended is True
        assert benchmark.was_actually_used is True

    def test_benchmark_to_redis_dict(self):
        """Test conversion to Redis dict."""
        benchmark = PerformanceBenchmark(
            operation="test_op",
            duration_seconds=1.0,
            success=True,
            metadata={"key": "value"},
        )

        redis_dict = benchmark.to_redis_dict()
        assert redis_dict["operation"] == "test_op"
        assert redis_dict["duration_seconds"] == 1.0
        assert redis_dict["success"] == "True"  # Booleans stored as strings in Redis
        assert isinstance(redis_dict["timestamp"], str)
        assert isinstance(redis_dict["metadata"], str)

    def test_benchmark_from_redis_dict(self):
        """Test reconstruction from Redis dict."""
        now = datetime.now(timezone.utc)
        redis_dict = {
            "timestamp": now.isoformat(),
            "operation": "test_op",
            "duration_seconds": 1.5,
            "success": "True",  # Booleans stored as strings in Redis
            "metadata": json.dumps({"test": "data"}),
            "recommendation_count": None,
            "confidence": None,
            "tool_name": None,
            "was_recommended": None,
            "was_actually_used": None,
            "conversation_id": None,
        }

        benchmark = PerformanceBenchmark.from_redis_dict(redis_dict)
        assert benchmark.operation == "test_op"
        assert benchmark.duration_seconds == 1.5
        assert benchmark.success is True  # Converted back to bool
        assert benchmark.metadata == {"test": "data"}


class TestBenchmarkStore:
    """Test BenchmarkStore functionality."""

    @pytest.fixture
    def mock_redis(self):
        """Create mock Redis client."""
        mock = AsyncMock()
        mock.hset = AsyncMock()
        mock.expire = AsyncMock()
        mock.zadd = AsyncMock()
        mock.zrevrangebyscore = AsyncMock(return_value=[])
        mock.hgetall = AsyncMock(return_value={})
        mock.aclose = AsyncMock()
        return mock

    @pytest.fixture
    def store(self, mock_redis):
        """Create benchmark store with mock Redis."""
        return BenchmarkStore(redis_client=mock_redis)

    @pytest.mark.asyncio
    async def test_record_benchmark(self, store, mock_redis):
        """Test recording a benchmark."""
        benchmark = PerformanceBenchmark(
            operation="test_op",
            duration_seconds=1.0,
            success=True,
        )

        await store.record(benchmark)

        # Verify Redis calls
        mock_redis.hset.assert_called_once()
        mock_redis.expire.assert_called()
        mock_redis.zadd.assert_called_once()

    @pytest.mark.asyncio
    async def test_record_benchmark_disabled(self, mock_redis):
        """Test recording when benchmarks are disabled."""
        with patch("src.core.benchmarks.config.ENABLE_BENCHMARKS", False):
            store = BenchmarkStore(redis_client=mock_redis)
            benchmark = PerformanceBenchmark(
                operation="test_op",
                duration_seconds=1.0,
                success=True,
            )

            await store.record(benchmark)

            # Should not call Redis
            mock_redis.hset.assert_not_called()

    @pytest.mark.asyncio
    async def test_record_benchmark_handles_errors(self, store, mock_redis):
        """Test recording handles Redis errors gracefully."""
        mock_redis.hset.side_effect = Exception("Redis error")

        benchmark = PerformanceBenchmark(
            operation="test_op",
            duration_seconds=1.0,
            success=True,
        )

        # Should not raise exception
        await store.record(benchmark)

    @pytest.mark.asyncio
    async def test_query_benchmarks(self, store, mock_redis):
        """Test querying benchmarks."""
        # Setup mock data
        now = datetime.now(timezone.utc)
        mock_key = f"benchmark:test_op:{int(now.timestamp() * 1000)}"
        mock_redis.zrevrangebyscore.return_value = [mock_key]

        # Mock hgetall to return proper data (booleans as strings, like Redis)
        mock_redis.hgetall.return_value = {
            "timestamp": now.isoformat(),
            "operation": "test_op",
            "duration_seconds": 1.5,  # Numeric, not string
            "success": "True",  # Booleans stored as strings in Redis
            "metadata": "{}",
            "recommendation_count": None,
            "confidence": None,
            "tool_name": None,
            "was_recommended": None,
            "was_actually_used": None,
            "conversation_id": None,
        }

        results = await store.query("test_op", limit=10)

        assert len(results) == 1
        assert results[0].operation == "test_op"
        mock_redis.zrevrangebyscore.assert_called_once()

    @pytest.mark.asyncio
    async def test_query_with_time_range(self, store, mock_redis):
        """Test querying with time range."""
        now = datetime.now(timezone.utc)
        start_time = now - timedelta(hours=1)
        end_time = now

        await store.query("test_op", start_time=start_time, end_time=end_time)

        # Verify time range was converted to timestamps
        call_args = mock_redis.zrevrangebyscore.call_args
        assert call_args is not None

    @pytest.mark.asyncio
    async def test_query_disabled_benchmarks(self, mock_redis):
        """Test querying when benchmarks are disabled."""
        with patch("src.core.benchmarks.config.ENABLE_BENCHMARKS", False):
            store = BenchmarkStore(redis_client=mock_redis)
            results = await store.query("test_op")
            assert results == []

    @pytest.mark.asyncio
    async def test_query_handles_errors(self, store, mock_redis):
        """Test query handles errors gracefully."""
        mock_redis.zrevrangebyscore.side_effect = Exception("Redis error")

        results = await store.query("test_op")
        assert results == []

    @pytest.mark.asyncio
    async def test_get_statistics(self, store, mock_redis):
        """Test getting statistics."""
        # Setup mock data with multiple benchmarks
        now = datetime.now(timezone.utc)
        mock_keys = [
            f"benchmark:test_op:{int((now - timedelta(seconds=i)).timestamp() * 1000)}"
            for i in range(3)
        ]
        mock_redis.zrevrangebyscore.return_value = mock_keys

        # Return different durations and success values
        benchmarks_data = [
            {"duration_seconds": "1.0", "success": "True"},
            {"duration_seconds": "2.0", "success": "True"},
            {"duration_seconds": "3.0", "success": "False"},
        ]

        async def mock_hgetall(key):
            idx = mock_keys.index(key)
            data = benchmarks_data[idx]
            return {
                "timestamp": now.isoformat(),
                "operation": "test_op",
                "duration_seconds": float(data["duration_seconds"]),
                "success": data["success"],  # Pass string through, from_redis_dict converts
                "metadata": "{}",
                "recommendation_count": None,
                "confidence": None,
                "tool_name": None,
                "was_recommended": None,
                "was_actually_used": None,
                "conversation_id": None,
            }

        mock_redis.hgetall.side_effect = mock_hgetall

        stats = await store.get_statistics("test_op")

        assert stats["count"] == 3
        assert stats["avg_duration"] == 2.0  # (1 + 2 + 3) / 3
        assert stats["min_duration"] == 1.0
        assert stats["max_duration"] == 3.0
        assert stats["success_rate"] == pytest.approx(66.67, rel=0.01)
        assert stats["total_successes"] == 2
        assert stats["total_failures"] == 1

    @pytest.mark.asyncio
    async def test_get_statistics_empty(self, store, mock_redis):
        """Test statistics with no data."""
        mock_redis.zrevrangebyscore.return_value = []

        stats = await store.get_statistics("test_op")

        assert stats["count"] == 0
        assert stats["avg_duration"] == 0.0
        assert stats["success_rate"] == 0.0

    @pytest.mark.asyncio
    async def test_get_tool_accuracy(self, store, mock_redis):
        """Test tool accuracy calculation."""
        # Setup mock data
        now = datetime.now(timezone.utc)
        mock_keys = [
            f"benchmark:tool_call:{int((now - timedelta(seconds=i)).timestamp() * 1000)}"
            for i in range(4)
        ]
        mock_redis.zrevrangebyscore.return_value = mock_keys

        # Different combinations of recommended/used
        tool_data = [
            {"was_recommended": "True", "was_actually_used": "True"},  # Good
            {"was_recommended": "True", "was_actually_used": "True"},  # Good
            {"was_recommended": "False", "was_actually_used": "True"},  # Missed
            {"was_recommended": "True", "was_actually_used": "False"},  # Not used
        ]

        async def mock_hgetall(key):
            idx = mock_keys.index(key)
            data = tool_data[idx]
            return {
                "timestamp": now.isoformat(),
                "operation": "tool_call",
                "duration_seconds": 1.0,
                "success": "True",  # Booleans stored as strings in Redis
                "metadata": "{}",
                "recommendation_count": None,
                "confidence": None,
                "tool_name": "test_tool",
                "conversation_id": None,
                "was_recommended": data["was_recommended"],  # Already strings
                "was_actually_used": data["was_actually_used"],  # Already strings
            }

        mock_redis.hgetall.side_effect = mock_hgetall

        accuracy = await store.get_tool_accuracy()

        assert accuracy["total_calls"] == 4
        assert accuracy["total_used"] == 3
        assert accuracy["recommended_and_used"] == 2
        assert accuracy["not_recommended_but_used"] == 1
        assert accuracy["precision"] == pytest.approx(66.67, rel=0.01)

    @pytest.mark.asyncio
    async def test_get_tool_accuracy_empty(self, store, mock_redis):
        """Test tool accuracy with no data."""
        mock_redis.zrevrangebyscore.return_value = []

        accuracy = await store.get_tool_accuracy()

        assert accuracy["total_calls"] == 0
        assert accuracy["precision"] == 0.0

    @pytest.mark.asyncio
    async def test_close(self, store, mock_redis):
        """Test closing the store."""
        await store.close()
        mock_redis.aclose.assert_called_once()

        # Client should be None after close
        assert store._client is None


class TestGlobalBenchmarkStore:
    """Test global benchmark store instance."""

    def test_get_benchmark_store(self):
        """Test getting global store instance."""
        store = get_benchmark_store()
        assert isinstance(store, BenchmarkStore)

    def test_get_benchmark_store_singleton(self):
        """Test store is singleton."""
        store1 = get_benchmark_store()
        store2 = get_benchmark_store()
        assert store1 is store2