tracer-cloud--opensre/tests/e2e/upstream_lambda/infrastructure_code/test_pipeline_infra.py

"""Infrastructure validation tests for upstream/downstream pipeline.

Tests that the AWS infrastructure and mock pipeline are working correctly.
"""

import json
import sys
import time
from datetime import datetime

import boto3
import requests


def get_stack_outputs(stack_name: str = "TracerUpstreamDownstreamTest") -> dict:
    """Get CloudFormation stack outputs."""
    cf_client = boto3.client("cloudformation")
    try:
        stack = cf_client.describe_stacks(StackName=stack_name)
        outputs = {o["OutputKey"]: o["OutputValue"] for o in stack["Stacks"][0]["Outputs"]}
        return outputs
    except Exception as e:
        print(f"ERROR: Could not get stack outputs: {e}")
        print(f"Make sure stack '{stack_name}' is deployed")
        sys.exit(1)


def test_happy_path(stack_outputs: dict) -> bool:
    """Test 1: Happy path with valid schema."""
    print("=" * 60)
    print("TEST 1: Happy Path")
    print("=" * 60)

    # Reset API to good schema
    try:
        requests.post(
            f"{stack_outputs['MockApiUrl']}/config",
            json={"inject_schema_change": False},
            timeout=10,
        )
        print("✓ Configured Mock API (good schema)")
    except Exception as e:
        print(f"✗ Failed to configure Mock API: {e}")
        return False

    # Verify API returns good data
    api_data = requests.get(f"{stack_outputs['MockApiUrl']}/data", timeout=10).json()
    has_customer_id = "customer_id" in api_data.get("data", [{}])[0]
    print(f"✓ Mock API returns data with customer_id: {has_customer_id}")
    if not has_customer_id:
        print("✗ TEST FAILED: Mock API should have customer_id")
        return False

    # Invoke ingester
    correlation_id = f"happy-{datetime.utcnow().strftime('%Y%m%d-%H%M%S')}"
    lambda_client = boto3.client("lambda")

    print(f"Invoking Ingester Lambda (correlation_id={correlation_id})...")
    response = lambda_client.invoke(
        FunctionName=stack_outputs["IngesterFunctionName"],
        Payload=json.dumps({"correlation_id": correlation_id}),
    )

    result = json.loads(response["Payload"].read())
    if result.get("statusCode") != 200:
        print(f"✗ Ingester failed: {result}")
        return False

    s3_key = result["s3_key"]
    print(f"✓ Data written to landing: {s3_key}")

    # Verify landing bucket has data
    s3_client = boto3.client("s3")
    try:
        obj = s3_client.get_object(Bucket=stack_outputs["LandingBucketName"], Key=s3_key)
        landing_data = json.loads(obj["Body"].read())
        has_customer_id = "customer_id" in landing_data.get("data", [{}])[0]
        print(f"✓ Landing data has customer_id: {has_customer_id}")

        if not has_customer_id:
            print("✗ TEST FAILED: Landing data should have customer_id")
            return False
    except Exception as e:
        print(f"✗ Failed to read landing data: {e}")
        return False

    # Wait for Mock DAG processing (S3 event trigger)
    print("Waiting 15s for Mock DAG to process...")
    time.sleep(15)

    # Verify processed output exists
    processed_key = s3_key.replace("ingested/", "processed/")
    try:
        obj = s3_client.get_object(Bucket=stack_outputs["ProcessedBucketName"], Key=processed_key)
        processed_data = json.loads(obj["Body"].read())

        # Verify transformation occurred
        has_amount_cents = "amount_cents" in processed_data.get("data", [{}])[0]
        print(f"✓ Processed output exists: {processed_key}")
        print(f"✓ Transformation added amount_cents: {has_amount_cents}")

        if not has_amount_cents:
            print("✗ TEST FAILED: Transform should add amount_cents")
            return False

        print("✓ HAPPY PATH PASSED")
        return True

    except s3_client.exceptions.NoSuchKey:
        print(f"✗ HAPPY PATH FAILED: No processed output at {processed_key}")
        return False
    except Exception as e:
        print(f"✗ HAPPY PATH FAILED: {e}")
        return False


def test_failure_path(stack_outputs: dict) -> bool:
    """Test 2: Failure path with schema mismatch."""
    print("\n" + "=" * 60)
    print("TEST 2: Failure Path")
    print("=" * 60)

    # Inject schema change
    try:
        requests.post(
            f"{stack_outputs['MockApiUrl']}/config",
            json={"inject_schema_change": True},
            timeout=10,
        )
        print("✓ Configured Mock API (schema change injected)")
    except Exception as e:
        print(f"✗ Failed to configure Mock API: {e}")
        return False

    # Verify API returns bad data
    api_data = requests.get(f"{stack_outputs['MockApiUrl']}/data", timeout=10).json()
    has_customer_id = "customer_id" in api_data.get("data", [{}])[0]
    print(f"✓ Mock API returns data WITHOUT customer_id: {not has_customer_id}")
    if has_customer_id:
        print("✗ TEST FAILED: Mock API should NOT have customer_id")
        return False

    # Invoke ingester
    correlation_id = f"fail-{datetime.utcnow().strftime('%Y%m%d-%H%M%S')}"
    lambda_client = boto3.client("lambda")

    print(f"Invoking Ingester Lambda (correlation_id={correlation_id})...")
    response = lambda_client.invoke(
        FunctionName=stack_outputs["IngesterFunctionName"],
        Payload=json.dumps({"correlation_id": correlation_id}),
    )

    result = json.loads(response["Payload"].read())
    s3_key = result["s3_key"]
    print(f"✓ Bad data written to landing: {s3_key}")

    # Verify landing bucket has bad data
    s3_client = boto3.client("s3")
    try:
        obj = s3_client.get_object(Bucket=stack_outputs["LandingBucketName"], Key=s3_key)
        landing_data = json.loads(obj["Body"].read())
        has_customer_id = "customer_id" in landing_data.get("data", [{}])[0]
        print(f"✓ Landing data missing customer_id: {not has_customer_id}")

        if has_customer_id:
            print("✗ TEST FAILED: Landing data should NOT have customer_id")
            return False
    except Exception as e:
        print(f"✗ Failed to read landing data: {e}")
        return False

    # Wait for Mock DAG processing
    print("Waiting 15s for Mock DAG to process and fail...")
    time.sleep(15)

    # Verify NO processed output (pipeline should have failed)
    processed_key = s3_key.replace("ingested/", "processed/")
    try:
        s3_client.head_object(Bucket=stack_outputs["ProcessedBucketName"], Key=processed_key)
        print(f"✗ FAILURE PATH FAILED: Output exists at {processed_key} (should have failed)")
        return False
    except Exception as e:
        # Should fail (object doesn't exist)
        if "404" in str(e) or "NoSuchKey" in str(e):
            print("✓ No processed output (correctly failed)")
        else:
            print(f"✗ Unexpected error checking processed output: {e}")
            return False

    # Verify error in CloudWatch logs
    logs_client = boto3.client("logs")
    log_group = f"/aws/lambda/{stack_outputs['MockDagFunctionName']}"

    try:
        response = logs_client.filter_log_events(
            logGroupName=log_group,
            startTime=int((time.time() - 300) * 1000),
            filterPattern="Schema validation failed",
        )

        if response["events"]:
            error_msg = response["events"][0]["message"]
            print("✓ Error logged in CloudWatch:")
            print(f"  {error_msg[:200]}")
            print("✓ FAILURE PATH PASSED")
            return True
        else:
            print("✗ FAILURE PATH FAILED: No error in logs")
            print(f"  Check log group: {log_group}")
            return False
    except Exception as e:
        print(f"✗ Failed to check CloudWatch logs: {e}")
        return False


if __name__ == "__main__":
    # Get stack outputs
    print("Getting stack outputs...")
    stack_outputs = get_stack_outputs()

    # Run tests
    happy_passed = test_happy_path(stack_outputs)
    failure_passed = test_failure_path(stack_outputs)

    # Summary
    print("\n" + "=" * 60)
    print("TEST SUMMARY")
    print("=" * 60)
    print(f"  Happy Path:   {'PASSED' if happy_passed else 'FAILED'}")
    print(f"  Failure Path: {'PASSED' if failure_passed else 'FAILED'}")
    print()

    if all([happy_passed, failure_passed]):
        print("✓ ALL TESTS PASSED")
        sys.exit(0)
    else:
        print("✗ SOME TESTS FAILED")
        sys.exit(1)