chopratejas--headroom/tests/test_integrations/agno/test_hooks.py

"""Tests for Agno hooks integration.

Tests cover:
1. HeadroomPreHook - Pre-hook for tracking before LLM calls
2. HeadroomPostHook - Post-hook for tracking after LLM calls
3. create_headroom_hooks - Convenience function
"""

from datetime import datetime
from unittest.mock import MagicMock

import pytest

# Check if Agno is available
try:
    import agno  # noqa: F401

    AGNO_AVAILABLE = True
except ImportError:
    AGNO_AVAILABLE = False

from headroom import HeadroomConfig, HeadroomMode

# Skip all tests if Agno not installed
pytestmark = pytest.mark.skipif(not AGNO_AVAILABLE, reason="Agno not installed")


class TestHeadroomPreHook:
    """Tests for HeadroomPreHook."""

    def test_init_defaults(self):
        """Initialize with default settings."""
        from headroom.integrations.agno import HeadroomPreHook

        hook = HeadroomPreHook()

        assert hook.mode == HeadroomMode.OPTIMIZE
        assert hook.model == "gpt-4o"
        assert hook.total_tokens_saved == 0
        assert hook.metrics_history == []

    def test_init_with_custom_config(self):
        """Initialize with custom config."""
        from headroom.integrations.agno import HeadroomPreHook

        config = HeadroomConfig(default_mode=HeadroomMode.AUDIT)
        hook = HeadroomPreHook(
            config=config,
            mode=HeadroomMode.SIMULATE,
            model="claude-3-5-sonnet-20241022",
        )

        assert hook.config is config
        assert hook.mode == HeadroomMode.SIMULATE
        assert hook.model == "claude-3-5-sonnet-20241022"

    def test_call_returns_input_unchanged(self):
        """Hook returns input unchanged (optimization at model level)."""
        from headroom.integrations.agno import HeadroomPreHook

        hook = HeadroomPreHook()

        run_input = "Hello, how are you?"
        result = hook(run_input)

        assert result == run_input

    def test_call_tracks_metrics(self):
        """Hook tracks metrics on each call."""
        from headroom.integrations.agno import HeadroomPreHook

        hook = HeadroomPreHook()

        hook("First input")
        hook("Second input")

        assert len(hook.metrics_history) == 2
        assert all(m.request_id for m in hook.metrics_history)
        assert all(isinstance(m.timestamp, datetime) for m in hook.metrics_history)

    def test_metrics_history_limited(self):
        """Metrics history is limited to 100 entries."""
        from headroom.integrations.agno import HeadroomPreHook

        hook = HeadroomPreHook()

        # Call 150 times
        for i in range(150):
            hook(f"Input {i}")

        assert len(hook.metrics_history) == 100

    def test_get_savings_summary_empty(self):
        """get_savings_summary with no history."""
        from headroom.integrations.agno import HeadroomPreHook

        hook = HeadroomPreHook()
        summary = hook.get_savings_summary()

        assert summary["total_requests"] == 0
        assert summary["total_tokens_saved"] == 0
        assert summary["average_savings_percent"] == 0

    def test_get_savings_summary_with_data(self):
        """get_savings_summary with metrics."""
        from headroom.integrations.agno import HeadroomPreHook

        hook = HeadroomPreHook()

        # Make some calls
        hook("Input 1")
        hook("Input 2")

        summary = hook.get_savings_summary()

        assert summary["total_requests"] == 2
        # Pre-hook doesn't do actual optimization, so tokens_saved is 0
        assert summary["total_tokens_saved"] == 0


class TestHeadroomPostHook:
    """Tests for HeadroomPostHook."""

    def test_init_defaults(self):
        """Initialize with default settings."""
        from headroom.integrations.agno import HeadroomPostHook

        hook = HeadroomPostHook()

        assert hook.log_level == "INFO"
        assert hook.token_alert_threshold is None
        assert hook.total_requests == 0
        assert hook.alerts == []

    def test_init_with_threshold(self):
        """Initialize with alert threshold."""
        from headroom.integrations.agno import HeadroomPostHook

        hook = HeadroomPostHook(
            log_level="DEBUG",
            token_alert_threshold=10000,
        )

        assert hook.log_level == "DEBUG"
        assert hook.token_alert_threshold == 10000

    def test_call_returns_output_unchanged(self):
        """Hook returns output unchanged."""
        from headroom.integrations.agno import HeadroomPostHook

        hook = HeadroomPostHook()

        output = MagicMock()
        output.content = "Hello!"
        result = hook(output)

        assert result is output

    def test_call_tracks_requests(self):
        """Hook tracks requests on each call."""
        from headroom.integrations.agno import HeadroomPostHook

        hook = HeadroomPostHook()

        output1 = MagicMock()
        output1.content = "First response"
        output2 = MagicMock()
        output2.content = "Second response"

        hook(output1)
        hook(output2)

        assert hook.total_requests == 2

    def test_call_extracts_token_metrics(self):
        """Hook extracts token metrics from response."""
        from headroom.integrations.agno import HeadroomPostHook

        hook = HeadroomPostHook()

        output = MagicMock()
        output.content = "Response"
        output.metrics = MagicMock()
        output.metrics.input_tokens = 50
        output.metrics.output_tokens = 20
        output.metrics.total_tokens = 70

        hook(output)

        assert hook._requests[0]["input_tokens"] == 50
        assert hook._requests[0]["output_tokens"] == 20
        assert hook._requests[0]["total_tokens"] == 70

    def test_call_triggers_alert(self):
        """Hook triggers alert when threshold exceeded."""
        from headroom.integrations.agno import HeadroomPostHook

        hook = HeadroomPostHook(token_alert_threshold=50)

        output = MagicMock()
        output.content = "Response"
        output.metrics = MagicMock()
        output.metrics.total_tokens = 100  # Exceeds threshold

        hook(output)

        assert len(hook.alerts) == 1
        assert "Token alert" in hook.alerts[0]
        assert "100" in hook.alerts[0]

    def test_call_no_alert_below_threshold(self):
        """No alert when tokens below threshold."""
        from headroom.integrations.agno import HeadroomPostHook

        hook = HeadroomPostHook(token_alert_threshold=100)

        output = MagicMock()
        output.content = "Response"
        output.metrics = MagicMock()
        output.metrics.total_tokens = 50  # Below threshold

        hook(output)

        assert len(hook.alerts) == 0

    def test_requests_limited(self):
        """Request history is limited to 1000 entries."""
        from headroom.integrations.agno import HeadroomPostHook

        hook = HeadroomPostHook()

        # Call 1500 times
        for i in range(1500):
            output = MagicMock()
            output.content = f"Response {i}"
            hook(output)

        assert len(hook._requests) == 1000

    def test_get_summary_empty(self):
        """get_summary with no requests."""
        from headroom.integrations.agno import HeadroomPostHook

        hook = HeadroomPostHook()
        summary = hook.get_summary()

        assert summary["total_requests"] == 0
        assert summary["total_tokens"] == 0
        assert summary["alerts"] == 0

    def test_get_summary_with_data(self):
        """get_summary with requests."""
        from headroom.integrations.agno import HeadroomPostHook

        hook = HeadroomPostHook()

        # Add some requests directly
        hook._requests = [
            {"total_tokens": 100},
            {"total_tokens": 200},
            {"total_tokens": 50},
        ]

        summary = hook.get_summary()

        assert summary["total_requests"] == 3
        assert summary["total_tokens"] == 350
        assert summary["average_tokens"] == 350 / 3

    def test_reset(self):
        """reset() clears all state."""
        from headroom.integrations.agno import HeadroomPostHook

        hook = HeadroomPostHook()

        # Add some state
        hook._requests = [{"test": 1}]
        hook._alerts = ["alert"]

        hook.reset()

        assert hook._requests == []
        assert hook._alerts == []


class TestCreateHeadroomHooks:
    """Tests for create_headroom_hooks convenience function."""

    def test_returns_tuple(self):
        """Returns tuple of (pre_hook, post_hook)."""
        from headroom.integrations.agno import (
            HeadroomPostHook,
            HeadroomPreHook,
            create_headroom_hooks,
        )

        pre_hook, post_hook = create_headroom_hooks()

        assert isinstance(pre_hook, HeadroomPreHook)
        assert isinstance(post_hook, HeadroomPostHook)

    def test_passes_config_to_pre_hook(self):
        """Passes config to pre_hook."""
        from headroom.integrations.agno import create_headroom_hooks

        config = HeadroomConfig(default_mode=HeadroomMode.AUDIT)
        pre_hook, _ = create_headroom_hooks(config=config)

        assert pre_hook.config is config

    def test_passes_mode_to_pre_hook(self):
        """Passes mode to pre_hook."""
        from headroom.integrations.agno import create_headroom_hooks

        pre_hook, _ = create_headroom_hooks(mode=HeadroomMode.SIMULATE)

        assert pre_hook.mode == HeadroomMode.SIMULATE

    def test_passes_model_to_pre_hook(self):
        """Passes model to pre_hook."""
        from headroom.integrations.agno import create_headroom_hooks

        pre_hook, _ = create_headroom_hooks(model="claude-3-5-sonnet-20241022")

        assert pre_hook.model == "claude-3-5-sonnet-20241022"

    def test_passes_log_level_to_post_hook(self):
        """Passes log_level to post_hook."""
        from headroom.integrations.agno import create_headroom_hooks

        _, post_hook = create_headroom_hooks(log_level="DEBUG")

        assert post_hook.log_level == "DEBUG"

    def test_passes_threshold_to_post_hook(self):
        """Passes token_alert_threshold to post_hook."""
        from headroom.integrations.agno import create_headroom_hooks

        _, post_hook = create_headroom_hooks(token_alert_threshold=5000)

        assert post_hook.token_alert_threshold == 5000

    def test_all_parameters(self):
        """Test with all parameters."""
        from headroom.integrations.agno import create_headroom_hooks

        config = HeadroomConfig()
        pre_hook, post_hook = create_headroom_hooks(
            config=config,
            mode=HeadroomMode.AUDIT,
            model="gpt-4-turbo",
            log_level="WARNING",
            token_alert_threshold=8000,
        )

        assert pre_hook.config is config
        assert pre_hook.mode == HeadroomMode.AUDIT
        assert pre_hook.model == "gpt-4-turbo"
        assert post_hook.log_level == "WARNING"
        assert post_hook.token_alert_threshold == 8000