trycua--cua/libs/python/agent/tests/test_telemetry_events.py

"""
Test script to verify telemetry events are emitted correctly.
"""

from unittest.mock import AsyncMock, MagicMock, patch

import pytest


class TestAgentTelemetryEvents:
    """Test telemetry events emitted by ComputerAgent."""

    @patch("cua_agent.agent.record_event")
    @patch("cua_agent.agent.is_telemetry_enabled", return_value=True)
    def test_agent_init_event(self, mock_telemetry_enabled, mock_record_event):
        """Test that agent_init event is emitted with correct args_provided."""
        from cua_agent.agent import ComputerAgent

        # Create agent with various args
        agent = ComputerAgent(
            model="anthropic/claude-sonnet-4-5-20250929",
            instructions="Test instructions",
            max_retries=5,  # non-default
            trajectory_dir="/tmp/test",
        )

        # Find the agent_init call
        agent_init_calls = [
            call for call in mock_record_event.call_args_list if call[0][0] == "agent_init"
        ]

        assert len(agent_init_calls) == 1, "agent_init should be called once"

        event_name, event_data = agent_init_calls[0][0]
        assert event_name == "agent_init"
        assert event_data["model"] == "anthropic/claude-sonnet-4-5-20250929"
        assert "instructions" in event_data["args_provided"]
        assert "max_retries" in event_data["args_provided"]
        assert "trajectory_dir" in event_data["args_provided"]

    @patch("cua_agent.agent.record_event")
    @patch("cua_agent.agent.is_telemetry_enabled", return_value=True)
    def test_agent_init_minimal_args(self, mock_telemetry_enabled, mock_record_event):
        """Test agent_init with minimal args (defaults)."""
        from cua_agent.agent import ComputerAgent

        agent = ComputerAgent(model="anthropic/claude-sonnet-4-5-20250929")

        agent_init_calls = [
            call for call in mock_record_event.call_args_list if call[0][0] == "agent_init"
        ]

        assert len(agent_init_calls) == 1
        event_name, event_data = agent_init_calls[0][0]

        # With defaults, only model-related things should be tracked
        # instructions, trajectory_dir, etc. should NOT be in args_provided
        assert "instructions" not in event_data["args_provided"]
        assert "trajectory_dir" not in event_data["args_provided"]
        assert "max_retries" not in event_data["args_provided"]  # default is 3

    @patch("cua_agent.agent.record_event")
    @patch("cua_agent.agent.is_telemetry_enabled", return_value=False)
    def test_no_events_when_telemetry_disabled(self, mock_telemetry_enabled, mock_record_event):
        """Test that no events are emitted when telemetry is disabled."""
        from cua_agent.agent import ComputerAgent

        agent = ComputerAgent(
            model="anthropic/claude-sonnet-4-5-20250929",
            telemetry_enabled=False,
        )

        # No agent_init should be called (telemetry disabled)
        agent_init_calls = [
            call for call in mock_record_event.call_args_list if call[0][0] == "agent_init"
        ]

        assert len(agent_init_calls) == 0


class TestActionTelemetryEvents:
    """Test telemetry events for computer actions."""

    @pytest.mark.asyncio
    @patch("cua_agent.agent.record_event")
    @patch("cua_agent.agent.is_telemetry_enabled", return_value=True)
    async def test_computer_action_executed_event(self, mock_telemetry_enabled, mock_record_event):
        """Test that computer_action_executed is emitted for computer calls."""
        from cua_agent.agent import ComputerAgent

        agent = ComputerAgent(model="anthropic/claude-sonnet-4-5-20250929")
        agent.telemetry_enabled = True

        # Mock computer handler
        mock_computer = MagicMock()
        mock_computer.click = AsyncMock(return_value=None)
        mock_computer.screenshot = AsyncMock(return_value="base64screenshot")

        # Create a mock computer_call item
        item = {
            "type": "computer_call",
            "call_id": "test-call-id",
            "action": {
                "type": "click",
                "x": 100,
                "y": 200,
            },
        }

        # Process the item (this would normally happen in the agent loop)
        # Note: We can't easily test this without running the full agent loop
        # This is more of an integration test

        # For unit testing, we verify the event structure
        expected_event = {
            "action_type": "click",
        }

        # Verify event structure is correct
        assert "action_type" in expected_event


class TestToolExecutedEvents:
    """Test telemetry events for tool execution."""

    def test_event_structure(self):
        """Test that agent_tool_executed event has correct structure."""
        expected_computer_tool_event = {
            "tool_type": "computer",
            "tool_name": "click",
        }

        expected_function_tool_event = {
            "tool_type": "function",
            "tool_name": "my_custom_function",
        }

        # Verify expected structure
        assert "tool_type" in expected_computer_tool_event
        assert "tool_name" in expected_computer_tool_event
        assert expected_computer_tool_event["tool_type"] in ["computer", "function"]


if __name__ == "__main__":
    pytest.main([__file__, "-v"])