chopratejas--headroom/tests/integrations/test_strands/test_hooks.py

"""Real-world integration tests for Strands HeadroomHookProvider.

These tests use actual AWS Bedrock API calls with real credentials.
NO MOCKS - all tests hit the real Bedrock API.

Skip in CI if AWS credentials are not available.
"""

from __future__ import annotations

import json
import os

import pytest

# Check for AWS credentials availability
SKIP_BEDROCK = not (
    os.environ.get("AWS_ACCESS_KEY_ID")
    or os.environ.get("AWS_PROFILE")
    or os.path.exists(os.path.expanduser("~/.aws/credentials"))
)

# Check if strands-agents is installed
try:
    from strands import Agent, tool
    from strands.models import BedrockModel

    STRANDS_AVAILABLE = True
except ImportError:
    STRANDS_AVAILABLE = False

    # Provide a no-op decorator when strands is not installed
    def tool(fn):
        return fn

    Agent = None  # type: ignore
    BedrockModel = None  # type: ignore

# Skip all tests if dependencies not available
pytestmark = [
    pytest.mark.skipif(SKIP_BEDROCK, reason="AWS credentials not available"),
    pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed"),
]


# ============================================================================
# Test Tools - Generate realistic verbose data for compression testing
# These are defined with @tool decorator for use when strands is installed.
# When strands is not installed, the no-op decorator ensures import succeeds.
# ============================================================================


@tool
def search_logs(query: str, limit: int = 100) -> str:
    """Search application logs. Returns JSON array of log entries.

    Args:
        query: Search query to find in logs
        limit: Maximum number of log entries to return

    Returns:
        JSON array of log entry objects
    """
    # Generate realistic verbose log data that should be compressed
    logs = [
        {
            "timestamp": f"2024-01-{(i % 28) + 1:02d}T{10 + (i % 12):02d}:00:00Z",
            "level": ["INFO", "DEBUG", "WARN", "ERROR"][i % 4],
            "service": ["api-gateway", "auth-service", "data-processor", "cache-service"][i % 4],
            "message": f"Request processed successfully - latency={50 + i}ms, query={query}",
            "request_id": f"req-{i:06d}-{hash(query) % 10000:04d}",
            "status_code": [200, 201, 400, 500][i % 4],
            "user_agent": "Mozilla/5.0 (compatible; TestBot/1.0)",
            "ip_address": f"192.168.{i % 256}.{(i * 7) % 256}",
            "trace_id": f"trace-{i:08x}",
            "span_id": f"span-{i:04x}",
            "duration_ms": 50 + (i * 3) % 200,
            "memory_mb": 128 + (i * 5) % 512,
            "cpu_percent": 10 + (i * 2) % 80,
        }
        for i in range(limit)
    ]
    return json.dumps(logs, indent=2)


@tool
def get_small_status() -> str:
    """Get a small status response that should NOT be compressed.

    Returns:
        Small JSON status object
    """
    return json.dumps({"status": "healthy", "uptime_seconds": 12345, "version": "1.2.3"})


@tool
def get_error_data() -> str:
    """Get error information. Error results should NOT be compressed.

    Returns:
        Error information (but not as a tool error)
    """
    return json.dumps(
        {
            "errors": [
                {"code": "E001", "message": "Connection timeout"},
                {"code": "E002", "message": "Authentication failed"},
            ],
            "timestamp": "2024-01-15T10:00:00Z",
        }
    )


@tool
def fetch_user_data(user_id: str) -> str:
    """Fetch detailed user data. Returns large JSON payload.

    Args:
        user_id: The user ID to fetch data for

    Returns:
        Large JSON object with user details
    """
    # Generate a large user profile that should trigger compression
    activities = [
        {
            "activity_id": f"act-{i:06d}",
            "type": ["login", "purchase", "view", "share"][i % 4],
            "timestamp": f"2024-01-{(i % 28) + 1:02d}T{10 + (i % 12):02d}:30:00Z",
            "details": {
                "ip": f"10.0.{i % 256}.{(i * 3) % 256}",
                "device": ["desktop", "mobile", "tablet"][i % 3],
                "browser": ["Chrome", "Firefox", "Safari"][i % 3],
                "duration_seconds": 30 + i * 5,
                "page_views": 1 + i % 10,
            },
            "metadata": {
                "session_id": f"sess-{i:08x}",
                "referrer": f"https://example.com/page/{i}",
                "utm_source": ["google", "facebook", "twitter", "email"][i % 4],
            },
        }
        for i in range(50)
    ]

    return json.dumps(
        {
            "user_id": user_id,
            "profile": {
                "name": "Test User",
                "email": f"{user_id}@example.com",
                "created_at": "2023-01-01T00:00:00Z",
            },
            "activities": activities,
        },
        indent=2,
    )


@tool
def simple_calculator(a: int, b: int, operation: str) -> str:
    """Simple calculator for basic operations.

    Args:
        a: First number
        b: Second number
        operation: One of 'add', 'subtract', 'multiply', 'divide'

    Returns:
        The result of the operation
    """
    if operation == "add":
        result = a + b
    elif operation == "subtract":
        result = a - b
    elif operation == "multiply":
        result = a * b
    elif operation == "divide":
        result = a / b if b != 0 else "undefined"
    else:
        result = "unknown operation"

    return json.dumps({"operation": operation, "a": a, "b": b, "result": result})


# ============================================================================
# Test Class
# ============================================================================


@pytest.mark.skipif(SKIP_BEDROCK, reason="AWS credentials not available")
@pytest.mark.skipif(not STRANDS_AVAILABLE, reason="strands-agents not installed")
class TestHeadroomHookProviderReal:
    """Real-world integration tests for HeadroomHookProvider with Bedrock."""

    @pytest.fixture
    def bedrock_model(self):
        """Create a BedrockModel instance using Claude 3 Haiku (fast and cheap)."""
        return BedrockModel(
            model_id="anthropic.claude-3-haiku-20240307-v1:0",
            region_name="us-west-2",
            temperature=0.1,  # Low temperature for consistent tests
        )

    @pytest.fixture
    def hook_provider(self):
        """Create a HeadroomHookProvider with test configuration."""
        from headroom.integrations.strands import HeadroomHookProvider

        return HeadroomHookProvider(
            compress_tool_outputs=True,
            min_tokens_to_compress=50,  # Low threshold for testing
            preserve_errors=True,
        )

    def test_hook_compresses_large_tool_output(self, bedrock_model, hook_provider):
        """Test that large tool outputs are compressed by the hook.

        This test:
        1. Creates an agent with the search_logs tool
        2. Asks a question that triggers the tool
        3. Verifies the hook compressed the output and saved tokens
        """
        # Create agent with hook provider
        agent = Agent(
            model=bedrock_model,
            tools=[search_logs],
            hooks=[hook_provider],
        )

        # Ask a question that will trigger the search_logs tool
        result = agent(
            "Search the logs for 'error' and tell me how many entries you found. "
            "Use limit=100 to get plenty of results."
        )

        # Verify the agent got a response
        assert result is not None

        # Check hook metrics
        metrics = hook_provider.get_savings_summary()

        # The hook should have processed at least one tool call
        assert metrics["total_requests"] >= 1, "Hook should have processed tool calls"

        # With 100 log entries, compression should have occurred
        # and saved significant tokens
        if metrics["compressed_requests"] > 0:
            assert metrics["total_tokens_saved"] > 0, "Should have saved tokens"
            assert metrics["total_tokens_before"] > metrics["total_tokens_after"]

    def test_hook_preserves_small_outputs(self, bedrock_model, hook_provider):
        """Test that small tool outputs are NOT compressed.

        This test:
        1. Creates an agent with a tool returning small output
        2. Triggers the tool
        3. Verifies the hook did not modify the small output
        """
        # Reset metrics from any previous tests
        hook_provider.reset()

        agent = Agent(
            model=bedrock_model,
            tools=[get_small_status],
            hooks=[hook_provider],
        )

        # Ask a question that will trigger the small status tool
        result = agent("What is the current system status? Use the get_small_status tool.")

        assert result is not None

        # Check metrics - small outputs should not be compressed
        metrics = hook_provider.get_savings_summary()

        # Tool was called but output was below threshold
        if metrics["total_requests"] > 0:
            # For small outputs, tokens_before == tokens_after (no compression)
            for m in hook_provider.metrics_history:
                if m.tool_name == "get_small_status" or "small" in str(m.skip_reason):
                    # Either not compressed or skip reason indicates below threshold
                    assert not m.was_compressed or m.skip_reason is not None, (
                        "Small output should not be compressed"
                    )

    def test_hook_preserves_errors(self, bedrock_model):
        """Test that error results are NOT compressed when preserve_errors=True.

        This test:
        1. Creates a hook with preserve_errors=True
        2. Creates an agent with a tool that returns error data
        3. Verifies error results are preserved unchanged
        """
        from headroom.integrations.strands import HeadroomHookProvider

        # Create hook with preserve_errors=True (default)
        hook_with_preserve = HeadroomHookProvider(
            compress_tool_outputs=True,
            min_tokens_to_compress=10,  # Very low threshold
            preserve_errors=True,
        )

        agent = Agent(
            model=bedrock_model,
            tools=[get_error_data],
            hooks=[hook_with_preserve],
        )

        # Get error data
        result = agent("Get the error data using get_error_data tool and summarize it.")

        assert result is not None

        # Check that error-related results were handled appropriately
        metrics = hook_with_preserve.get_savings_summary()

        # The get_error_data tool returns data about errors but doesn't itself error
        # So it should be processed normally (this tests the flow works)
        assert metrics["total_requests"] >= 0  # May or may not have been called

    def test_hook_metrics_tracking(self, bedrock_model, hook_provider):
        """Test that metrics are tracked correctly across multiple tool calls.

        This test:
        1. Creates an agent with multiple tools
        2. Makes requests that trigger various tools
        3. Verifies metrics are accumulated correctly
        """
        # Reset metrics
        hook_provider.reset()

        agent = Agent(
            model=bedrock_model,
            tools=[search_logs, get_small_status, simple_calculator],
            hooks=[hook_provider],
        )

        # First request - should trigger search_logs (large output)
        agent("Search logs for 'test' with limit=50 and give me a count.")

        # Second request - should trigger calculator (small output)
        agent("Calculate 15 + 27 using the calculator tool.")

        # Third request - should trigger status (small output)
        agent("Get the system status using get_small_status.")

        # Check accumulated metrics
        metrics = hook_provider.get_savings_summary()

        # Should have tracked multiple requests
        assert metrics["total_requests"] >= 1, "Should have tracked tool requests"

        # total_tokens_before should be >= total_tokens_after
        assert metrics["total_tokens_before"] >= metrics["total_tokens_after"]

        # History should contain records
        history = hook_provider.metrics_history
        assert len(history) >= 1, "Should have metrics history entries"

        # Each metric should have required fields
        for m in history:
            assert m.request_id is not None
            assert m.timestamp is not None
            assert m.tokens_before >= 0
            assert m.tokens_after >= 0

    def test_multiple_tool_calls_in_single_request(self, bedrock_model, hook_provider):
        """Test that multiple tool calls in a single agent request are all processed.

        This test:
        1. Asks a complex question requiring multiple tools
        2. Verifies each tool call is processed by the hook
        """
        # Reset metrics
        hook_provider.reset()

        agent = Agent(
            model=bedrock_model,
            tools=[search_logs, simple_calculator, fetch_user_data],
            hooks=[hook_provider],
        )

        # Ask a complex question that might trigger multiple tools
        result = agent(
            "I need you to do three things: "
            "1. Search logs for 'api' with limit=30. "
            "2. Calculate 100 * 5 using the calculator. "
            "3. Tell me the total number of results from step 1."
        )

        assert result is not None

        # Check that multiple tool calls were processed
        metrics = hook_provider.get_savings_summary()

        # Should have processed at least the search_logs call
        assert metrics["total_requests"] >= 1

        # Verify metrics history
        history = hook_provider.metrics_history

        # At minimum, should have processed search_logs (which has large output)
        # The actual tools called depend on the model's interpretation
        assert len(history) >= 1

        # Check that we have tool names recorded
        tool_names = [m.tool_name for m in history]
        assert all(name is not None for name in tool_names)

    def test_hook_reset_clears_metrics(self, bedrock_model, hook_provider):
        """Test that reset() clears all accumulated metrics.

        This test:
        1. Makes some requests to accumulate metrics
        2. Calls reset()
        3. Verifies all metrics are cleared
        """
        agent = Agent(
            model=bedrock_model,
            tools=[search_logs],
            hooks=[hook_provider],
        )

        # Make a request to accumulate metrics
        agent("Search logs for 'test' with limit=20.")

        # Verify we have some metrics
        assert hook_provider.total_tokens_saved >= 0

        # Reset
        hook_provider.reset()

        # Verify metrics are cleared
        assert hook_provider.total_tokens_saved == 0
        assert len(hook_provider.metrics_history) == 0

        metrics = hook_provider.get_savings_summary()
        assert metrics["total_requests"] == 0
        assert metrics["total_tokens_saved"] == 0

    def test_hook_with_compression_disabled(self, bedrock_model):
        """Test that hook passes through without compression when disabled.

        This test:
        1. Creates a hook with compress_tool_outputs=False
        2. Verifies tool outputs are not modified
        """
        from headroom.integrations.strands import HeadroomHookProvider

        # Create hook with compression disabled
        disabled_hook = HeadroomHookProvider(
            compress_tool_outputs=False,
            min_tokens_to_compress=10,
        )

        agent = Agent(
            model=bedrock_model,
            tools=[search_logs],
            hooks=[disabled_hook],
        )

        result = agent("Search logs for 'api' with limit=50.")

        assert result is not None

        # When compression is disabled, no requests should be tracked
        # (the hook doesn't register callbacks when disabled)
        metrics = disabled_hook.get_savings_summary()
        assert metrics["compressed_requests"] == 0

    def test_hook_concurrent_safety(self, bedrock_model, hook_provider):
        """Test that hook is thread-safe for concurrent access.

        This test verifies that metrics tracking is thread-safe
        by checking that accumulated values are consistent.
        """
        import threading

        # Reset metrics
        hook_provider.reset()

        agent = Agent(
            model=bedrock_model,
            tools=[simple_calculator],
            hooks=[hook_provider],
        )

        results = []
        errors = []

        def make_request(n: int):
            try:
                result = agent(f"Calculate {n} + {n} using simple_calculator.")
                results.append(result)
            except Exception as e:
                errors.append(e)

        # Run a few sequential requests (concurrent Bedrock calls might be rate-limited)
        threads = []
        for i in range(3):
            t = threading.Thread(target=make_request, args=(i,))
            threads.append(t)
            t.start()
            # Small delay to avoid rate limiting
            import time

            time.sleep(0.5)

        for t in threads:
            t.join(timeout=60)  # 60 second timeout per thread

        # Check we got results (some may have failed due to rate limits)
        assert len(results) > 0 or len(errors) > 0

        # Metrics should still be consistent
        metrics = hook_provider.get_savings_summary()
        assert metrics["total_tokens_before"] >= metrics["total_tokens_after"]

    def test_hook_handles_empty_tool_response(self, bedrock_model, hook_provider):
        """Test that hook handles tools returning empty responses gracefully."""

        @tool
        def empty_response() -> str:
            """Return an empty response."""
            return ""

        hook_provider.reset()

        agent = Agent(
            model=bedrock_model,
            tools=[empty_response],
            hooks=[hook_provider],
        )

        # This might not trigger the tool if the model decides it's not needed
        result = agent("Call the empty_response tool and tell me what you got.")

        assert result is not None

        # Should handle gracefully without errors
        metrics = hook_provider.get_savings_summary()
        # Just verify no exceptions and metrics are valid
        assert metrics["total_tokens_before"] >= 0
        assert metrics["total_tokens_after"] >= 0