Implement Lambda response transformation to API Gateway format (#1927)

GarrettBeatty · web-flow · commit 1d5ce5493f82 · 2025-01-15T12:33:41.000-05:00
diff --git a/Tools/LambdaTestTool-v2/src/Amazon.Lambda.TestTool/Amazon.Lambda.TestTool.csproj b/Tools/LambdaTestTool-v2/src/Amazon.Lambda.TestTool/Amazon.Lambda.TestTool.csproj
@@ -1,4 +1,4 @@
-﻿<Project Sdk="Microsoft.NET.Sdk.Web">
+<Project Sdk="Microsoft.NET.Sdk.Web">
 
   <PropertyGroup>
 	  <Description>A tool to help debug and test your .NET AWS Lambda functions locally.</Description>
@@ -15,6 +15,7 @@
   </PropertyGroup>
 
   <ItemGroup>
+    <PackageReference Include="AWSSDK.Lambda" Version="3.7.408.1" />
     <PackageReference Include="Spectre.Console" Version="0.49.1" />
     <PackageReference Include="Spectre.Console.Cli" Version="0.49.1" />
     <PackageReference Include="Amazon.Lambda.APIGatewayEvents" Version="2.7.1" />
diff --git a/Tools/LambdaTestTool-v2/src/Amazon.Lambda.TestTool/Extensions/InvokeResponseExtensions.cs b/Tools/LambdaTestTool-v2/src/Amazon.Lambda.TestTool/Extensions/InvokeResponseExtensions.cs
@@ -0,0 +1,171 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+using System.Text.Json;
+using Amazon.Lambda.APIGatewayEvents;
+using Amazon.Lambda.Model;
+using Amazon.Lambda.TestTool.Models;
+
+/// <summary>
+/// Provides extension methods for converting Lambda InvokeResponse to API Gateway response types.
+/// </summary>
+public static class InvokeResponseExtensions
+{
+    /// <summary>
+    /// Converts an Amazon Lambda InvokeResponse to an APIGatewayProxyResponse.
+    /// </summary>
+    /// <param name="invokeResponse">The InvokeResponse from a Lambda function invocation.</param>
+    /// <param name="emulatorMode">The API Gateway emulator mode (Rest or Http).</param>
+    /// <returns>An APIGatewayProxyResponse object.</returns>
+    /// <remarks>
+    /// If the response cannot be deserialized as an APIGatewayProxyResponse, it returns an error response.
+    /// The error response differs based on the emulator mode:
+    /// - For Rest mode: StatusCode 502 with a generic error message.
+    /// - For Http mode: StatusCode 500 with a generic error message.
+    /// </remarks>
+    public static APIGatewayProxyResponse ToApiGatewayProxyResponse(this InvokeResponse invokeResponse, ApiGatewayEmulatorMode emulatorMode)
+    {
+        if (emulatorMode == ApiGatewayEmulatorMode.HttpV2)
+        {
+            throw new NotSupportedException("This function should only be used with Rest and Httpv1 emulator modes");
+        }
+
+        using var reader = new StreamReader(invokeResponse.Payload);
+        string responseJson = reader.ReadToEnd();
+        try
+        {
+            return JsonSerializer.Deserialize<APIGatewayProxyResponse>(responseJson);
+        }
+        catch
+        {
+            if (emulatorMode == ApiGatewayEmulatorMode.Rest)
+            {
+                return new APIGatewayProxyResponse
+                {
+                    StatusCode = 502,
+                    Body = "{\"message\":\"Internal server error\"}",
+                    Headers = new Dictionary<string, string>
+                {
+                    { "Content-Type", "application/json" }
+                },
+                    IsBase64Encoded = false
+                };
+            }
+            else
+            {
+                return new APIGatewayProxyResponse
+                {
+                    StatusCode = 500,
+                    Body = "{\"message\":\"Internal Server Error\"}",
+                    Headers = new Dictionary<string, string>
+                {
+                    { "Content-Type", "application/json" }
+                },
+                    IsBase64Encoded = false
+                };
+            }
+        }
+    }
+
+    /// <summary>
+    /// Converts an Amazon Lambda InvokeResponse to an APIGatewayHttpApiV2ProxyResponse.
+    /// </summary>
+    /// <param name="invokeResponse">The InvokeResponse from a Lambda function invocation.</param>
+    /// <returns>An APIGatewayHttpApiV2ProxyResponse object.</returns>
+    /// <remarks>
+    /// This method reads the payload from the InvokeResponse and passes it to ToHttpApiV2Response
+    /// for further processing and conversion.
+    /// </remarks>
+    public static APIGatewayHttpApiV2ProxyResponse ToApiGatewayHttpApiV2ProxyResponse(this InvokeResponse invokeResponse)
+    {
+        using var reader = new StreamReader(invokeResponse.Payload);
+        string responseJson = reader.ReadToEnd();
+        return ToHttpApiV2Response(responseJson);
+    }
+
+    /// <summary>
+    /// Converts a response string to an APIGatewayHttpApiV2ProxyResponse.
+    /// </summary>
+    /// <param name="response">The response string to convert.</param>
+    /// <returns>An APIGatewayHttpApiV2ProxyResponse object.</returns>
+    /// <remarks>
+    /// This method replicates the observed behavior of API Gateway's HTTP API 
+    /// with Lambda integrations using payload format version 2.0, which differs 
+    /// from the official documentation.
+    /// 
+    /// Observed behavior:
+    /// 1. If the response is a JSON object with a 'statusCode' property:
+    ///    - It attempts to deserialize it as a full APIGatewayHttpApiV2ProxyResponse.
+    ///    - If deserialization fails, it returns a 500 Internal Server Error.
+    /// 2. For any other response (including non-JSON strings, invalid JSON, or partial JSON):
+    ///    - Sets statusCode to 200
+    ///    - Uses the response as-is for the body
+    ///    - Sets Content-Type to application/json
+    ///    - Sets isBase64Encoded to false
+    /// 
+    /// This behavior contradicts the official documentation, which states:
+    /// "If your Lambda function returns valid JSON and doesn't return a statusCode, 
+    /// API Gateway assumes a 200 status code and treats the entire response as the body."
+    /// 
+    /// In practice, API Gateway does not validate the JSON. It treats any response 
+    /// without a 'statusCode' property as a raw body, regardless of whether it's 
+    /// valid JSON or not.
+    /// 
+    /// For example, if a Lambda function returns:
+    ///     '{"name": "John Doe", "age":'
+    /// API Gateway will treat this as a raw string body in a 200 OK response, not attempting
+    /// to parse or validate the JSON structure.
+    /// 
+    /// This method replicates this observed behavior rather than the documented behavior.
+    /// </remarks>
+    private static APIGatewayHttpApiV2ProxyResponse ToHttpApiV2Response(string response)
+    {
+        try
+        {
+            // Try to deserialize as JsonElement first to inspect the structure
+            var jsonElement = JsonSerializer.Deserialize<JsonElement>(response);
+
+            // Check if it's an object that might represent a full response
+            if (jsonElement.ValueKind == JsonValueKind.Object &&
+                jsonElement.TryGetProperty("statusCode", out _))
+            {
+                // It has a statusCode property, so try to deserialize as full response
+                try
+                {
+                    return JsonSerializer.Deserialize<APIGatewayHttpApiV2ProxyResponse>(response);
+                }
+                catch
+                {
+                    // If deserialization fails, return Internal Server Error
+                    return new APIGatewayHttpApiV2ProxyResponse
+                    {
+                        StatusCode = 500,
+                        Body = "{\"message\":\"Internal Server Error\"}",
+                        Headers = new Dictionary<string, string>
+                    {
+                        { "Content-Type", "application/json" }
+                    },
+                        IsBase64Encoded = false
+                    };
+                }
+            }
+        }
+        catch
+        {
+            // If JSON parsing fails, fall through to default behavior
+        }
+
+        // Default behavior: return the response as-is
+        return new APIGatewayHttpApiV2ProxyResponse
+        {
+            StatusCode = 200,
+            Body = response,
+            Headers = new Dictionary<string, string>
+        {
+            { "Content-Type", "application/json" }
+        },
+            IsBase64Encoded = false
+        };
+    }
+
+}
diff --git a/Tools/LambdaTestTool-v2/tests/Amazon.Lambda.TestTool.IntegrationTests/InvokeResponseExtensionsIntegrationTests.cs b/Tools/LambdaTestTool-v2/tests/Amazon.Lambda.TestTool.IntegrationTests/InvokeResponseExtensionsIntegrationTests.cs
@@ -0,0 +1,196 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+using Amazon.Lambda.APIGatewayEvents;
+using Amazon.Lambda.Model;
+using Amazon.Lambda.TestTool.Models;
+using System.Text;
+using System.Text.Json;
+
+namespace Amazon.Lambda.TestTool.IntegrationTests;
+
+/// <summary>
+/// Integration tests for InvokeResponseExtensions.
+/// </summary>
+/// <remarks>
+/// Developer's Note:
+/// These tests don't have direct access to the intermediate result of the Lambda to API Gateway conversion.
+/// Instead, we test the final API Gateway response object to ensure our conversion methods produce results
+/// that match the actual API Gateway behavior. This approach allows us to verify the correctness of our
+/// conversion methods within the constraints of not having access to AWS's internal conversion process.
+/// </remarks>
+[Collection("ApiGateway Integration Tests")]
+public class InvokeResponseExtensionsIntegrationTests
+{
+    private readonly ApiGatewayIntegrationTestFixture _fixture;
+
+    public InvokeResponseExtensionsIntegrationTests(ApiGatewayIntegrationTestFixture fixture)
+    {
+        _fixture = fixture;
+    }
+
+    [Theory]
+    [InlineData(ApiGatewayEmulatorMode.Rest)]
+    [InlineData(ApiGatewayEmulatorMode.HttpV1)]
+    public async Task ToApiGatewayProxyResponse_ValidResponse_MatchesDirectConversion(ApiGatewayEmulatorMode emulatorMode)
+    {
+        // Arrange
+        var testResponse = new APIGatewayProxyResponse
+        {
+            StatusCode = 200,
+            Body = JsonSerializer.Serialize(new { message = "Hello, World!" }),
+            Headers = new Dictionary<string, string> { { "Content-Type", "application/json" } }
+        };
+        var invokeResponse = new InvokeResponse
+        {
+            Payload = new MemoryStream(Encoding.UTF8.GetBytes(JsonSerializer.Serialize(testResponse)))
+        };
+
+        // Act
+        var convertedResponse = invokeResponse.ToApiGatewayProxyResponse(emulatorMode);
+
+        // Assert
+        var apiUrl = emulatorMode == ApiGatewayEmulatorMode.Rest
+            ? _fixture.ParseAndReturnBodyRestApiUrl
+            : _fixture.ParseAndReturnBodyHttpApiV1Url;
+        var (actualResponse, httpTestResponse) = await _fixture.ApiGatewayTestHelper.ExecuteTestRequest(convertedResponse, apiUrl, emulatorMode);
+        await _fixture.ApiGatewayTestHelper.AssertResponsesEqual(actualResponse, httpTestResponse);
+    }
+
+    [Fact]
+    public async Task ToApiGatewayHttpApiV2ProxyResponse_ValidResponse_MatchesDirectConversion()
+    {
+        // Arrange
+        var testResponse = new APIGatewayHttpApiV2ProxyResponse
+        {
+            StatusCode = 200,
+            Body = JsonSerializer.Serialize(new { message = "Hello, World!" }),
+            Headers = new Dictionary<string, string> { { "Content-Type", "application/json" } }
+        };
+        var invokeResponse = new InvokeResponse
+        {
+            Payload = new MemoryStream(Encoding.UTF8.GetBytes(JsonSerializer.Serialize(testResponse)))
+        };
+
+        // Act
+        var convertedResponse = invokeResponse.ToApiGatewayHttpApiV2ProxyResponse();
+
+        // Assert
+        var (actualResponse, httpTestResponse) = await _fixture.ApiGatewayTestHelper.ExecuteTestRequest(convertedResponse, _fixture.ParseAndReturnBodyHttpApiV2Url);
+        await _fixture.ApiGatewayTestHelper.AssertResponsesEqual(actualResponse, httpTestResponse);
+    }
+
+    [Theory]
+    [InlineData(ApiGatewayEmulatorMode.Rest, 502, "Internal server error")]
+    [InlineData(ApiGatewayEmulatorMode.HttpV1, 500, "Internal Server Error")]
+    public async Task ToApiGatewayProxyResponse_InvalidJson_ReturnsErrorResponse(ApiGatewayEmulatorMode emulatorMode, int expectedStatusCode, string expectedErrorMessage)
+    {
+        // Arrange
+        var invokeResponse = new InvokeResponse
+        {
+            Payload = new MemoryStream(Encoding.UTF8.GetBytes("Not a valid proxy response object"))
+        };
+
+        // Act
+        var convertedResponse = invokeResponse.ToApiGatewayProxyResponse(emulatorMode);
+
+        // Assert
+        Assert.Equal(expectedStatusCode, convertedResponse.StatusCode);
+        Assert.Contains(expectedErrorMessage, convertedResponse.Body);
+
+        var apiUrl = emulatorMode == ApiGatewayEmulatorMode.Rest
+            ? _fixture.ParseAndReturnBodyRestApiUrl
+            : _fixture.ParseAndReturnBodyHttpApiV1Url;
+        var (actualResponse, _) = await _fixture.ApiGatewayTestHelper.ExecuteTestRequest(convertedResponse, apiUrl, emulatorMode);
+        Assert.Equal(expectedStatusCode, (int)actualResponse.StatusCode);
+        var content = await actualResponse.Content.ReadAsStringAsync();
+        Assert.Contains(expectedErrorMessage, content);
+    }
+
+    /// <summary>
+    /// Tests various Lambda return values to verify API Gateway's handling of responses.
+    /// </summary>
+    /// <param name="responsePayload">The payload returned by the Lambda function.</param>
+    /// <remarks>
+    /// This test demonstrates a discrepancy between the official AWS documentation
+    /// and the actual observed behavior of API Gateway HTTP API v2 with Lambda
+    /// proxy integrations (payload format version 2.0).
+    /// 
+    /// Official documentation states:
+    /// "If your Lambda function returns valid JSON and doesn't return a statusCode,
+    /// API Gateway assumes a 200 status code and treats the entire response as the body."
+    /// 
+    /// However, the observed behavior (which this test verifies) is:
+    /// - API Gateway does not validate whether the returned data is valid JSON.
+    /// - Any response from the Lambda function that is not a properly formatted
+    ///   API Gateway response object (i.e., an object with a 'statusCode' property)
+    ///   is treated as a raw body in a 200 OK response.
+    /// - This includes valid JSON objects without a statusCode, JSON arrays, 
+    ///   primitive values, and invalid JSON strings.
+    /// 
+    /// This test ensures that our ToApiGatewayHttpApiV2ProxyResponse method
+    /// correctly replicates this observed behavior, rather than the documented behavior.
+    /// </remarks>
+    [Theory]
+    [InlineData("{\"name\": \"John Doe\", \"age\":")]  // Invalid JSON (partial object)
+    [InlineData("{\"name\": \"John Doe\", \"age\": 30}")]  // Valid JSON object without statusCode
+    [InlineData("[1, 2, 3, 4, 5]")]  // JSON array
+    [InlineData("\"Hello, World!\"")]  // String primitive
+    [InlineData("42")]  // Number primitive
+    [InlineData("true")]  // Boolean primitive
+    public async Task ToApiGatewayHttpApiV2ProxyResponse_VariousPayloads_ReturnsAsRawBody(string responsePayload)
+    {
+        // Arrange
+        var invokeResponse = new InvokeResponse
+        {
+            Payload = new MemoryStream(Encoding.UTF8.GetBytes(responsePayload))
+        };
+
+        // Act
+        var convertedResponse = invokeResponse.ToApiGatewayHttpApiV2ProxyResponse();
+
+        // Assert
+        Assert.Equal(200, convertedResponse.StatusCode);
+        Assert.Equal(responsePayload, convertedResponse.Body);
+        Assert.Equal("application/json", convertedResponse.Headers["Content-Type"]);
+
+        // Verify against actual API Gateway behavior
+        var (actualResponse, httpTestResponse) = await _fixture.ApiGatewayTestHelper.ExecuteTestRequest(convertedResponse, _fixture.ParseAndReturnBodyHttpApiV2Url);
+        await _fixture.ApiGatewayTestHelper.AssertResponsesEqual(actualResponse, httpTestResponse);
+
+        // Additional checks for API Gateway specific behavior
+        Assert.Equal(200, (int)actualResponse.StatusCode);
+        var content = await actualResponse.Content.ReadAsStringAsync();
+        Assert.Equal(responsePayload, content);
+        Assert.Equal("application/json", actualResponse.Content.Headers.ContentType?.ToString());
+    }
+
+    [Fact]
+    public async Task ToApiGatewayHttpApiV2ProxyResponse_StatusCodeAsFloat_ReturnsInternalServerError()
+    {
+        // Arrange
+        var responsePayload = "{\"statusCode\": 200.5, \"body\": \"Hello\", \"headers\": {\"Content-Type\": \"text/plain\"}}";
+        var invokeResponse = new InvokeResponse
+        {
+            Payload = new MemoryStream(Encoding.UTF8.GetBytes(responsePayload))
+        };
+
+        // Act
+        var convertedResponse = invokeResponse.ToApiGatewayHttpApiV2ProxyResponse();
+
+        // Assert
+        Assert.Equal(500, convertedResponse.StatusCode);
+        Assert.Equal("{\"message\":\"Internal Server Error\"}", convertedResponse.Body);
+        Assert.Equal("application/json", convertedResponse.Headers["Content-Type"]);
+
+        // Verify against actual API Gateway behavior
+        var (actualResponse, httpTestResponse) = await _fixture.ApiGatewayTestHelper.ExecuteTestRequest(convertedResponse, _fixture.ParseAndReturnBodyHttpApiV2Url);
+        await _fixture.ApiGatewayTestHelper.AssertResponsesEqual(actualResponse, httpTestResponse);
+
+        // Additional checks for API Gateway specific behavior
+        Assert.Equal(500, (int)actualResponse.StatusCode);
+        var content = await actualResponse.Content.ReadAsStringAsync();
+        Assert.Equal("{\"message\":\"Internal Server Error\"}", content);
+        Assert.Equal("application/json", actualResponse.Content.Headers.ContentType?.ToString());
+    }
+}
diff --git a/Tools/LambdaTestTool-v2/tests/Amazon.Lambda.TestTool.UnitTests/Extensions/InvokeResponseExtensionsTests.cs b/Tools/LambdaTestTool-v2/tests/Amazon.Lambda.TestTool.UnitTests/Extensions/InvokeResponseExtensionsTests.cs
