Update header logic and add test cases

Tools/LambdaTestTool-v2/src/Amazon.Lambda.TestTool/Extensions/ApiGatewayResponseExtensions.cs

@@ -6,89 +6,63 @@
 using System;
 using System.IO;
 using System.Text;
+using System.Text.Json;
 
-/// <summary>
-/// Provides extension methods for converting API Gateway responses to HttpResponse objects.
-/// </summary>
 public static class ApiGatewayResponseExtensions
 {
-    /// <summary>
-    /// Converts an APIGatewayProxyResponse to an HttpResponse.
-    /// </summary>
-    /// <param name="apiResponse">The API Gateway proxy response to convert.</param>
-    /// <returns>An HttpResponse object representing the API Gateway response.</returns>
     public static HttpResponse ToHttpResponse(this APIGatewayProxyResponse apiResponse)
     {
         var httpContext = new DefaultHttpContext();
         var response = httpContext.Response;
 
-        response.StatusCode = apiResponse.StatusCode;
-
         SetResponseHeaders(response, apiResponse.Headers, apiResponse.MultiValueHeaders);
-
         SetResponseBody(response, apiResponse.Body, apiResponse.IsBase64Encoded);
+        SetContentTypeAndStatusCode(response, apiResponse.Headers, apiResponse.MultiValueHeaders, apiResponse.StatusCode, isV2: false);
 
         return response;
     }
 
-    /// <summary>
-    /// Converts an APIGatewayHttpApiV2ProxyResponse to an HttpResponse.
-    /// </summary>
-    /// <param name="apiResponse">The API Gateway HTTP API v2 proxy response to convert.</param>
-    /// <returns>An HttpResponse object representing the API Gateway response.</returns>
     public static HttpResponse ToHttpResponse(this APIGatewayHttpApiV2ProxyResponse apiResponse)
     {
         var httpContext = new DefaultHttpContext();
         var response = httpContext.Response;
 
-        response.StatusCode = apiResponse.StatusCode;
-
         SetResponseHeaders(response, apiResponse.Headers);
-
-        if (apiResponse.Cookies != null)
-        {
-            foreach (var cookie in apiResponse.Cookies)
-            {
-                response.Headers.Append("Set-Cookie", cookie);
-            }
-        }
-
         SetResponseBody(response, apiResponse.Body, apiResponse.IsBase64Encoded);
+        SetContentTypeAndStatusCode(response, apiResponse.Headers, null, apiResponse.StatusCode, isV2: true);
 
         return response;
     }
 
-    /// <summary>
-    /// Sets the headers on the HttpResponse object.
-    /// </summary>
-    /// <param name="response">The HttpResponse object to modify.</param>
-    /// <param name="headers">The single-value headers to set.</param>
-    /// <param name="multiValueHeaders">The multi-value headers to set.</param>
     private static void SetResponseHeaders(HttpResponse response, IDictionary<string, string>? headers, IDictionary<string, IList<string>>? multiValueHeaders = null)
     {
-        if (headers != null)
+        var processedHeaders = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
+
+        if (multiValueHeaders != null)
         {
-            foreach (var header in headers)
+            foreach (var header in multiValueHeaders)
             {
-                response.Headers[header.Key] = header.Value;
+                response.Headers[header.Key] = new StringValues([.. header.Value]);
+                processedHeaders.Add(header.Key);
             }
         }
 
-        if (multiValueHeaders != null)
+        if (headers != null)
         {
-            foreach (var header in multiValueHeaders)
+            foreach (var header in headers)
             {
-                response.Headers[header.Key] = new StringValues(header.Value.ToArray());
+                if (!processedHeaders.Contains(header.Key))
+                {
+                    response.Headers[header.Key] = header.Value;
+                }
+                else
+                {
+                    response.Headers.Append(header.Key, header.Value);
+                }
             }
         }
     }
 
-    /// <summary>
-    /// Sets the body of the HttpResponse object.
-    /// </summary>
-    /// <param name="response">The HttpResponse object to modify.</param>
-    /// <param name="body">The body content to set.</param>
-    /// <param name="isBase64Encoded">Indicates whether the body is Base64 encoded.</param>
     private static void SetResponseBody(HttpResponse response, string? body, bool isBase64Encoded)
     {
         if (!string.IsNullOrEmpty(body))
@@ -103,6 +77,70 @@ private static void SetResponseBody(HttpResponse response, string? body, bool is
                 bodyBytes = Encoding.UTF8.GetBytes(body);
             }
             response.Body = new MemoryStream(bodyBytes);
+            response.ContentLength = bodyBytes.Length;
+        }
+    }
+
+    private static void SetContentTypeAndStatusCode(HttpResponse response, IDictionary<string, string>? headers, IDictionary<string, IList<string>>? multiValueHeaders, int statusCode, bool isV2)
+    {
+        string? contentType = null;
+
+        if (headers != null && headers.TryGetValue("Content-Type", out var headerContentType))
+        {
+            contentType = headerContentType;
+        }
+        else if (multiValueHeaders != null && multiValueHeaders.TryGetValue("Content-Type", out var multiValueContentType))
+        {
+            contentType = multiValueContentType[0];
+        }
+
+        bool isValidJson = false;
+        if (isV2 && contentType == null && response.Body != null)
+        {
+            response.Body.Seek(0, SeekOrigin.Begin);
+            using var reader = new StreamReader(response.Body, leaveOpen: true);
+            var bodyContent = reader.ReadToEnd();
+            response.Body.Seek(0, SeekOrigin.Begin);
+
+            isValidJson = IsValidJson(bodyContent);
+            if (isValidJson)
+            {
+                contentType = "application/json";
+            }
+        }
+
+        if (contentType != null)
+        {
+            response.ContentType = contentType;
+        }
+
+        if (isV2 && isValidJson && statusCode == 0)
+        {
+            response.StatusCode = 200;
+        }
+        else
+        {
+            response.StatusCode = statusCode;
+        }
+    }
+
+    private static bool IsValidJson(string strInput)
+    {
+        if (string.IsNullOrWhiteSpace(strInput)) { return false; }
+        strInput = strInput.Trim();
+        if ((strInput.StartsWith("{") && strInput.EndsWith("}")) ||
+            (strInput.StartsWith("[") && strInput.EndsWith("]")))
+        {
+            try
+            {
+                var obj = JsonSerializer.Deserialize<object>(strInput);
+                return true;
+            }
+            catch (JsonException)
+            {
+                return false;
+            }
         }
+        return false;
     }
 }

Tools/LambdaTestTool-v2/src/Amazon.Lambda.TestTool/Extensions/ServiceCollectionExtensions.cs

@@ -1,31 +1,31 @@
-using Amazon.Lambda.TestTool.Services;
-using Amazon.Lambda.TestTool.Services.IO;
-using Microsoft.Extensions.DependencyInjection.Extensions;
-
-namespace Amazon.Lambda.TestTool.Extensions;
-
-/// <summary>
-/// A class that contains extension methods for the <see cref="IServiceCollection"/> interface.
-/// </summary>
-public static class ServiceCollectionExtensions
-{
+using Amazon.Lambda.TestTool.Services;
+using Amazon.Lambda.TestTool.Services.IO;
+using Microsoft.Extensions.DependencyInjection.Extensions;
+
+namespace Amazon.Lambda.TestTool.Extensions;
+
+/// <summary>
+/// A class that contains extension methods for the <see cref="IServiceCollection"/> interface.
+/// </summary>
+public static class ServiceCollectionExtensions
+{
     /// <summary>
     /// Adds a set of services for the .NET CLI portion of this application.
-    /// </summary>
-    public static void AddCustomServices(this IServiceCollection serviceCollection,
-        ServiceLifetime lifetime = ServiceLifetime.Singleton)
-    {
-        serviceCollection.TryAdd(new ServiceDescriptor(typeof(IToolInteractiveService), typeof(ConsoleInteractiveService), lifetime));
-        serviceCollection.TryAdd(new ServiceDescriptor(typeof(IDirectoryManager), typeof(DirectoryManager), lifetime));
+    /// </summary>
+    public static void AddCustomServices(this IServiceCollection serviceCollection,
+        ServiceLifetime lifetime = ServiceLifetime.Singleton)
+    {
+        serviceCollection.TryAdd(new ServiceDescriptor(typeof(IToolInteractiveService), typeof(ConsoleInteractiveService), lifetime));
+        serviceCollection.TryAdd(new ServiceDescriptor(typeof(IDirectoryManager), typeof(DirectoryManager), lifetime));
     }
 
     /// <summary>
     /// Adds a set of services for the API Gateway emulator portion of this application.
-    /// </summary>
-    public static void AddApiGatewayEmulatorServices(this IServiceCollection serviceCollection,
-        ServiceLifetime lifetime = ServiceLifetime.Singleton)
-    {
-        serviceCollection.TryAdd(new ServiceDescriptor(typeof(IApiGatewayRouteConfigService), typeof(ApiGatewayRouteConfigService), lifetime));
-        serviceCollection.TryAdd(new ServiceDescriptor(typeof(IEnvironmentManager), typeof(EnvironmentManager), lifetime));
-    }
+    /// </summary>
+    public static void AddApiGatewayEmulatorServices(this IServiceCollection serviceCollection,
+        ServiceLifetime lifetime = ServiceLifetime.Singleton)
+    {
+        serviceCollection.TryAdd(new ServiceDescriptor(typeof(IApiGatewayRouteConfigService), typeof(ApiGatewayRouteConfigService), lifetime));
+        serviceCollection.TryAdd(new ServiceDescriptor(typeof(IEnvironmentManager), typeof(EnvironmentManager), lifetime));
+    }
 }