tracer-cloud--opensre/tests/e2e/upstream_lambda/test_agent_e2e.py

#!/usr/bin/env python3
"""End-to-end agent investigation test for upstream/downstream pipeline.

Triggers a failure in the pipeline and tests if the agent can correctly investigate and diagnose it.
"""

import json
import sys
import time
from datetime import UTC, datetime

import boto3
import requests

from platform.observability.trace.hook import traceable
from surfaces.cli.investigation import run_investigation_cli
from tests.utils.alert_factory import create_alert
from tests.utils.conftest import UPSTREAM_DOWNSTREAM_CONFIG


def _pick_field(payload: dict, keys: list[str]) -> str | None:
    for key in keys:
        value = payload.get(key)
        if isinstance(value, str) and value:
            return value
    data = payload.get("data")
    if isinstance(data, dict):
        for key in keys:
            value = data.get(key)
            if isinstance(value, str) and value:
                return value
    return None


def _parse_ingester_payload(response: requests.Response) -> dict:
    try:
        result = response.json()
    except ValueError as exc:
        raise RuntimeError(f"Ingester API did not return JSON: {response.text}") from exc

    payload: object = result
    if isinstance(result, dict) and "statusCode" in result and "body" in result:
        status_code = result.get("statusCode")
        if status_code != 200:
            raise RuntimeError(f"Pipeline trigger failed: {result}")
        payload = result.get("body")
        if isinstance(payload, str):
            try:
                payload = json.loads(payload)
            except json.JSONDecodeError as exc:
                raise RuntimeError(f"Ingester body was not valid JSON: {payload}") from exc

    if not isinstance(payload, dict):
        raise RuntimeError(f"Unexpected ingester payload: {payload}")
    if payload.get("error"):
        raise RuntimeError(f"Pipeline trigger failed: {payload}")

    s3_key = _pick_field(payload, ["s3_key", "s3Key", "key"])
    bucket = _pick_field(payload, ["s3_bucket", "s3Bucket", "bucket"])
    if not s3_key or not bucket:
        raise RuntimeError(f"Ingester response missing s3 fields: {payload}")

    return {"s3_key": s3_key, "bucket": bucket, "payload": payload}


def trigger_pipeline_failure() -> dict:
    """Trigger a pipeline failure and return alert data."""
    print("=" * 60)
    print("Triggering Pipeline Failure")
    print("=" * 60)

    # Trigger failure via HTTP
    correlation_id = f"alert-local-{datetime.now(UTC).strftime('%Y%m%d-%H%M%S')}"

    print(f"\nTriggering pipeline with schema error (correlation_id={correlation_id})...")
    response = requests.post(
        UPSTREAM_DOWNSTREAM_CONFIG["ingester_api_url"],
        json={"correlation_id": correlation_id, "inject_schema_change": True},
        timeout=10,
    )

    parsed = _parse_ingester_payload(response)
    s3_key = parsed["s3_key"]
    bucket = parsed["bucket"]

    print(f"✓ Bad data written to: s3://{bucket}/{s3_key}")
    print("Waiting 10s for Mock DAG to process and fail...")
    time.sleep(10)

    # Get error from CloudWatch logs
    logs_client = boto3.client("logs")
    log_group = f"/aws/lambda/{UPSTREAM_DOWNSTREAM_CONFIG['mock_dag_function_name']}"

    print(f"Checking logs in: {log_group}")
    response = logs_client.filter_log_events(
        logGroupName=log_group,
        startTime=int((time.time() - 120) * 1000),
        filterPattern=correlation_id,
    )

    error_message = "Schema validation failed"
    for event in response["events"]:
        if "PIPELINE FAILED" in event["message"]:
            error_message = event["message"].split("Error: ")[-1].split("\n")[0]
            break

    print(f"✓ Error detected: {error_message}")

    return {
        "correlation_id": correlation_id,
        "s3_key": s3_key,
        "bucket": bucket,
        "error_message": error_message,
        "log_group": log_group,
    }


def test_agent_investigation(failure_data: dict) -> bool:
    """Test agent can investigate the pipeline failure."""
    print("\n" + "=" * 60)
    print("Testing Agent Investigation")
    print("=" * 60)

    pipeline_name = "upstream_downstream_pipeline"
    run_id = f"run_{datetime.now(UTC).strftime('%Y%m%d_%H%M%S')}"

    # Create alert
    raw_alert = create_alert(
        pipeline_name=pipeline_name,
        run_name=run_id,
        status="failed",
        timestamp=datetime.now(UTC).isoformat(),
        annotations={
            "s3_bucket": failure_data["bucket"],
            "s3_key": failure_data["s3_key"],
            "correlation_id": failure_data["correlation_id"],
            "error": failure_data["error_message"],
            "lambda_log_group": failure_data["log_group"],
            "function_name": UPSTREAM_DOWNSTREAM_CONFIG["mock_dag_function_name"],
            "landing_bucket": UPSTREAM_DOWNSTREAM_CONFIG["landing_bucket_name"],
            "processed_bucket": UPSTREAM_DOWNSTREAM_CONFIG["processed_bucket_name"],
            "mock_api_url": UPSTREAM_DOWNSTREAM_CONFIG["mock_api_url"],
            "ingester_function": UPSTREAM_DOWNSTREAM_CONFIG["ingester_function_name"],
            "mock_dag_function": UPSTREAM_DOWNSTREAM_CONFIG["mock_dag_function_name"],
            "context_sources": "s3,lambda,cloudwatch",
        },
    )

    print("\nAlert created:")
    print(f"  Alert ID: {raw_alert['alert_id']}")
    print(f"  Pipeline: {pipeline_name}")
    print(f"  S3 Key: {failure_data['s3_key']}")
    print(f"  Correlation ID: {failure_data['correlation_id']}")

    print("\nRunning agent investigation...")

    @traceable(
        run_type="chain",
        name=f"test_lambda_upstream - {raw_alert['alert_id'][:8]}",
        metadata={
            "alert_id": raw_alert["alert_id"],
            "pipeline_name": pipeline_name,
            "correlation_id": failure_data["correlation_id"],
            "s3_key": failure_data["s3_key"],
            "lambda_function": failure_data.get("mock_dag_function"),
        },
    )
    def run_investigation():
        return run_investigation_cli(raw_alert=raw_alert)

    try:
        result = run_investigation()

        print("\n✓ Investigation complete")
        print(f"  Root cause: {result.get('root_cause', 'Unknown')}")

        return True

    except Exception as e:
        print(f"\n✗ Investigation failed: {e}")
        import traceback

        traceback.print_exc()
        return False


def main():
    """Main test flow."""
    print("=" * 60)
    print("Upstream/Downstream Pipeline - Agent E2E Test")
    print("=" * 60)
    print()

    # Step 1: Trigger failure
    failure_data = trigger_pipeline_failure()

    # Step 2: Test agent investigation
    success = test_agent_investigation(failure_data)

    if success:
        print("\n" + "=" * 60)
        print("✓ AGENT E2E TEST PASSED")
        print("=" * 60)
        return 0
    else:
        print("\n" + "=" * 60)
        print("✗ AGENT E2E TEST FAILED")
        print("=" * 60)
        return 1


if __name__ == "__main__":
    sys.exit(main())