Add more tests

Author: normj

Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.IntegrationTests/ApiGatewayStreamingTests.cs

@@ -127,6 +127,91 @@ public async Task StreamingErrorEndpoint_StreamIsTruncated()
                 _output.WriteLine($"Expected error: {ex.Message}");
             }
         }
+
+        [Fact]
+        public async Task OnCompletedCallback_IsExecuted()
+        {
+            var apiUrl = await _fixture.GetApiUrlAsync();
+            using var httpClient = new HttpClient();
+
+            // First request registers the OnCompleted callback
+            var response = await httpClient.GetWithRetryAsync($"{apiUrl}oncompleted-test");
+            Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+            var body = await response.Content.ReadAsStringAsync();
+            _output.WriteLine($"Body: {body}");
+            Assert.Contains("OnCompleted callback registered", body);
+
+            // Second request verifies the callback was executed
+            var verifyResponse = await httpClient.GetWithRetryAsync($"{apiUrl}oncompleted-verify");
+            Assert.Equal(HttpStatusCode.OK, verifyResponse.StatusCode);
+            var verifyBody = await verifyResponse.Content.ReadAsStringAsync();
+            _output.WriteLine($"Verify body: {verifyBody}");
+
+            var doc = JsonDocument.Parse(verifyBody);
+            Assert.True(doc.RootElement.GetProperty("onCompletedExecuted").GetBoolean(),
+                "OnCompleted callback should have been executed");
+        }
+
+        [Fact]
+        public async Task CustomHeaders_PassedThroughApiGateway()
+        {
+            var apiUrl = await _fixture.GetApiUrlAsync();
+            using var httpClient = new HttpClient();
+
+            var response = await httpClient.GetWithRetryAsync($"{apiUrl}custom-headers", HttpStatusCode.Created);
+
+            _output.WriteLine($"Status: {response.StatusCode}");
+            Assert.Equal(HttpStatusCode.Created, response.StatusCode);
+
+            var body = await response.Content.ReadAsStringAsync();
+            Assert.Contains("Custom headers response", body);
+
+            // Verify custom headers are present in the response
+            Assert.True(response.Headers.Contains("X-Custom-Header"), "X-Custom-Header should be present");
+            Assert.Equal("custom-value", response.Headers.GetValues("X-Custom-Header").First());
+            Assert.True(response.Headers.Contains("X-Another-Header"), "X-Another-Header should be present");
+            Assert.Equal("another-value", response.Headers.GetValues("X-Another-Header").First());
+        }
+
+        [Fact]
+        public async Task SetCookie_PassedThroughApiGateway()
+        {
+            var apiUrl = await _fixture.GetApiUrlAsync();
+            var handler = new HttpClientHandler { UseCookies = false };
+            using var httpClient = new HttpClient(handler);
+
+            var response = await httpClient.GetWithRetryAsync($"{apiUrl}set-cookie");
+
+            Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+
+            var body = await response.Content.ReadAsStringAsync();
+            _output.WriteLine($"Body: {body}");
+            Assert.Contains("Cookies set", body);
+
+            // Verify Set-Cookie headers are present
+            Assert.True(response.Headers.Contains("Set-Cookie"), "Set-Cookie header should be present");
+            var cookies = response.Headers.GetValues("Set-Cookie").ToList();
+            _output.WriteLine($"Cookies: {string.Join("; ", cookies)}");
+            Assert.True(cookies.Any(c => c.Contains("session=abc123")), "session cookie should be present");
+            Assert.True(cookies.Any(c => c.Contains("theme=dark")), "theme cookie should be present");
+        }
+
+        [Fact]
+        public async Task PostWithBody_EchoesRequestBody()
+        {
+            var apiUrl = await _fixture.GetApiUrlAsync();
+            using var httpClient = new HttpClient();
+
+            var content = new StringContent("Hello from integration test", Encoding.UTF8, "text/plain");
+            var response = await httpClient.PostWithRetryAsync($"{apiUrl}echo-body", content);
+
+            _output.WriteLine($"Status: {response.StatusCode}");
+            var body = await response.Content.ReadAsStringAsync();
+            _output.WriteLine($"Body: {body}");
+
+            Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+            Assert.Contains("Echo: Hello from integration test", body);
+        }
     }
 
     internal static class HttpClientExtension
@@ -151,6 +236,27 @@ public static async Task<HttpResponseMessage> GetWithRetryAsync(this HttpClient
             }
             throw new Exception($"Failed to get expected status code {expectedCode} from {url} after {maxRetries} attempts");
         }
+
+        public static async Task<HttpResponseMessage> PostWithRetryAsync(this HttpClient httpClient, string url, HttpContent content, HttpStatusCode expectedCode = HttpStatusCode.OK, int maxRetries = 5, int delaySeconds = 5)
+        {
+            for (var i = 0; i < maxRetries; i++)
+            {
+                try
+                {
+                    var response = await httpClient.PostAsync(url, content);
+                    if (response.StatusCode == expectedCode)
+                    {
+                        return response;
+                    }
+                }
+                catch
+                {
+                    // Ignore and retry
+                }
+                await Task.Delay(TimeSpan.FromSeconds(delaySeconds));
+            }
+            throw new Exception($"Failed to get expected status code {expectedCode} from {url} after {maxRetries} attempts");
+        }
     }
 
     /// <summary>

Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/AspNetCoreStreamingApiGatewayTest/Program.cs

@@ -53,4 +53,73 @@
     return Results.Json(new { message = "Hello from streaming Lambda", timestamp = DateTime.UtcNow.ToString("o") });
 });
 
+app.MapGet("/oncompleted-test", async (HttpContext context) =>
+{
+    // Register an OnCompleted callback that writes a marker to a response header.
+    // Since headers are sent in the prelude before the body, we use a different approach:
+    // write a marker into the body from the OnCompleted callback via a shared flag.
+    var completedMarker = new CompletedMarker();
+    context.Response.RegisterForDispose(completedMarker);
+
+    context.Response.OnCompleted(async (state) =>
+    {
+        var marker = (CompletedMarker)state;
+        marker.WasExecuted = true;
+        // Write to a static so the next request can verify it ran
+        CompletedMarkerStore.LastMarkerExecuted = true;
+    }, completedMarker);
+
+    context.Response.ContentType = "text/plain";
+    context.Response.StatusCode = 200;
+
+    var stream = context.Response.BodyWriter.AsStream();
+    using var writer = new StreamWriter(stream, leaveOpen: true);
+    await writer.WriteAsync("OnCompleted callback registered");
+    await writer.FlushAsync();
+});
+
+app.MapGet("/oncompleted-verify", (HttpContext context) =>
+{
+    // Returns whether the OnCompleted callback from the previous request was executed
+    return Results.Json(new { onCompletedExecuted = CompletedMarkerStore.LastMarkerExecuted });
+});
+
+app.MapGet("/custom-headers", (HttpContext context) =>
+{
+    context.Response.StatusCode = 201;
+    context.Response.ContentType = "text/plain";
+    context.Response.Headers["X-Custom-Header"] = "custom-value";
+    context.Response.Headers["X-Another-Header"] = "another-value";
+    return Results.Text("Custom headers response", "text/plain", statusCode: 201);
+});
+
+app.MapGet("/set-cookie", (HttpContext context) =>
+{
+    context.Response.Cookies.Append("session", "abc123", new CookieOptions
+    {
+        Path = "/",
+        HttpOnly = true
+    });
+    context.Response.Cookies.Append("theme", "dark");
+    return Results.Text("Cookies set");
+});
+
+app.MapPost("/echo-body", async (HttpContext context) =>
+{
+    using var reader = new StreamReader(context.Request.Body);
+    var body = await reader.ReadToEndAsync();
+    return Results.Text($"Echo: {body}");
+});
+
 app.Run();
+
+class CompletedMarker : IDisposable
+{
+    public bool WasExecuted { get; set; }
+    public void Dispose() { }
+}
+
+static class CompletedMarkerStore
+{
+    public static bool LastMarkerExecuted { get; set; }
+}