test_lambda_handler.py

import json
import unittest
from unittest.mock import Mock, patch

import lambda_handler as handler


class LambdaHandlerTests(unittest.TestCase):
    def test_transform_session_logs_uses_v1_style_entries(self):
        session = {
            "trace": [
                {
                    "type": "api_start",
                    "timestamp": 1712500000000,
                    "payload": {
                        "url": "https://example.com/api",
                        "method": "POST",
                        "headers": {"Authorization": "Bearer secret", "Content-Type": "application/json"},
                        "body": {"phone": "5551234567"},
                    },
                },
                {
                    "type": "api_end",
                    "timestamp": 1712500001000,
                    "payload": {
                        "url": "https://example.com/api",
                        "status": 200,
                        "response": {"ok": True},
                        "duration_ms": 1000,
                    },
                },
                {
                    "type": "tool_start",
                    "timestamp": 1712500002000,
                    "payload": {
                        "name": "transfer_call",
                        "arguments": {"department": "parts"},
                        "trace_id": "trace-1",
                    },
                },
                {
                    "type": "tool_end",
                    "timestamp": 1712500003000,
                    "payload": {"name": "transfer_call", "success": True, "trace_id": "trace-1"},
                },
            ]
        }

        result = handler.transform_session_logs(session)

        self.assertEqual(
            [entry["event"] for entry in result],
            ["external_api_call", "external_api_response", "tool_call_execution", "tool_call_result"],
        )
        self.assertEqual(result[0]["details"]["headers"]["Authorization"], "***")
        self.assertTrue(result[0]["details"]["has_auth"])
        self.assertEqual(result[2]["details"]["tool_name"], "transfer_call")
        self.assertEqual(result[3]["details"]["status"], "success")

    def test_build_api_logs_appends_transfer_event_when_trace_exists(self):
        body = {
            "end_timestamp": 1712500003000,
            "log": {
                "sessionLogData": {
                    "trace": [
                        {
                            "type": "tool_start",
                            "timestamp": 1712500002000,
                            "payload": {"name": "transfer_call", "arguments": {"department": "service"}},
                        }
                    ]
                },
                "data": {
                    "transferEventFromEngine": {
                        "transfer_to": "+15555550123",
                        "x-api-key": "should-not-leak",
                    }
                },
            },
        }

        result = handler.build_api_logs(body)

        self.assertIsInstance(result, list)
        self.assertEqual([entry["event"] for entry in result], ["tool_call_execution", "transfer_event"])
        self.assertEqual(result[1]["details"]["x-api-key"], "***")

    def test_build_api_logs_preserves_transfer_event_at_cap_boundary(self):
        trace = [
            {
                "type": "print",
                "timestamp": 1712500000000 + index,
                "payload": {"output": f"line {index}"},
            }
            for index in range(handler.MAX_API_LOG_ENTRIES)
        ]
        body = {
            "end_timestamp": "1712500999000",
            "log": {
                "sessionLogData": {"trace": trace},
                "data": {"transferEventFromEngine": {"transfer_to": "+15555550123"}},
            },
        }

        result = handler.build_api_logs(body)

        self.assertEqual(len(result), handler.MAX_API_LOG_ENTRIES)
        self.assertEqual(result[-1]["event"], "transfer_event")
        self.assertEqual(result[-1]["details"]["transfer_to"], "+15555550123")

    def test_build_api_logs_preserves_transfer_fallback_shape(self):
        body = {
            "log": {
                "data": {
                    "transferEventFromEngine": {
                        "transfer_to": "+15555550123",
                        "source": "engine",
                    }
                }
            }
        }

        result = handler.build_api_logs(body)

        self.assertEqual(result, {"transferEvent": {"transfer_to": "+15555550123", "source": "engine"}})

    def test_numeric_string_timestamp_is_formatted(self):
        self.assertEqual(handler._format_timestamp("1712500000000"), "2024-04-07T14:26:40+00:00")

    def test_unparseable_numeric_string_timestamp_falls_back(self):
        huge_timestamp = "9" * 100
        self.assertEqual(handler._format_timestamp(huge_timestamp), huge_timestamp)

    def test_action_start_uses_payload_or_inferred_method(self):
        explicit = handler._convert_trace_entry(
            "action_start",
            {"action": "api2.custom", "method": "PATCH", "url": "https://example.com"},
            "ts",
        )
        inferred = handler._convert_trace_entry(
            "action_start",
            {"action": "api2_put", "url": "https://example.com"},
            "ts",
        )

        self.assertEqual(explicit[0]["details"]["method"], "PATCH")
        self.assertEqual(inferred[0]["details"]["method"], "PUT")

    def test_list_truncation_includes_marker(self):
        result = handler._truncate_large_values(list(range(52)))

        self.assertEqual(len(result), 51)
        self.assertEqual(result[-1], "...[2 items truncated]")

    def test_post_req_uses_configured_timeout(self):
        response = Mock()
        response.ok = True
        response.status_code = 200
        response.json.return_value = {"success": True}

        with patch("lambda_handler.requests.post", return_value=response) as post:
            result = handler.post_req("https://example.com/webhook", {"hello": "world"})

        post.assert_called_once_with(
            "https://example.com/webhook",
            json={"hello": "world"},
            headers=None,
            timeout=handler.HTTP_TIMEOUT_SECONDS,
        )
        self.assertEqual(result, {"ok": True, "status_code": 200, "body": {"success": True}})

    def test_lambda_requires_import_api_key_before_posting(self):
        old_key = handler.DGA_IMPORT_API_KEY
        handler.DGA_IMPORT_API_KEY = ""
        try:
            with patch("lambda_handler.post_req") as post:
                result = handler.lambda_handler({"body": json.dumps({"log_id": "log_test"})}, None)
        finally:
            handler.DGA_IMPORT_API_KEY = old_key

        self.assertEqual(result["statusCode"], 500)
        post.assert_not_called()


if __name__ == "__main__":
    unittest.main()