Add additional test for asp.net workloads

Author: ppittle

Libraries/test/Amazon.Lambda.AspNetCoreServer.Test/Amazon.Lambda.AspNetCoreServer.Test.csproj

@@ -41,6 +41,7 @@
   <ItemGroup>
     <ProjectReference Include="..\..\src\Amazon.Lambda.APIGatewayEvents\Amazon.Lambda.APIGatewayEvents.csproj" />
     <ProjectReference Include="..\..\src\Amazon.Lambda.ApplicationLoadBalancerEvents\Amazon.Lambda.ApplicationLoadBalancerEvents.csproj" />
+    <ProjectReference Include="..\..\src\Amazon.Lambda.RuntimeSupport\Amazon.Lambda.RuntimeSupport.csproj" />
     <ProjectReference Include="..\TestWebApp\TestWebApp.csproj" />
     <ProjectReference Include="..\..\src\Amazon.Lambda.TestUtilities\Amazon.Lambda.TestUtilities.csproj" />
   </ItemGroup>

Libraries/test/Amazon.Lambda.AspNetCoreServer.Test/TestApiGatewayHttpApiV2Calls.cs

@@ -1,4 +1,4 @@
-using System;
+using System;
 using System.Collections.Generic;
 using System.IO;
 using System.IO.Compression;
@@ -9,13 +9,15 @@
 using System.Threading.Tasks;
 
 using Amazon.Lambda.APIGatewayEvents;
+using Amazon.Lambda.RuntimeSupport;
+using Amazon.Lambda.Serialization.SystemTextJson;
 using Amazon.Lambda.TestUtilities;
 
 using Newtonsoft.Json;
 using Newtonsoft.Json.Linq;
 
 using TestWebApp;
-
+using TestWebApp.Controllers;
 using Xunit;
 
 
@@ -282,6 +284,30 @@ public async Task TestTraceIdSetFromLambdaContext()
             }
         }
 
+        #if NET8_0_OR_GREATER
+        /// <summary>
+        /// Verifies that <see cref="HttpV2LambdaFunction.GetBeforeSnapshotRequests"/> is invoked during startup.
+        /// </summary>
+        /// <returns></returns>
+        [Fact]
+        public async Task TestSnapStartInitialization()
+        {
+            using var e1 = new EnvironmentVariableHelper("AWS_LAMBDA_FUNCTION_NAME", nameof(TestSnapStartInitialization));
+            using var e2 = new EnvironmentVariableHelper("AWS_LAMBDA_INITIALIZATION_TYPE", "snap-start");
+            using var e3 = new EnvironmentVariableHelper("AWS_LAMBDA_RUNTIME_API", "localhost:123");
+
+            var t = LambdaBootstrapBuilder.Create<APIGatewayHttpApiV2ProxyRequest>(
+                    new TestWebApp.HttpV2LambdaFunction().FunctionHandlerAsync,
+                    new DefaultLambdaJsonSerializer())
+                .Build()
+                .RunAsync();
+
+            await Task.Delay(1000);
+
+            Assert.True(SnapStartController.Invoked);
+        }
+        #endif
+
         private async Task<APIGatewayHttpApiV2ProxyResponse> InvokeAPIGatewayRequest(string fileName, bool configureApiToReturnExceptionDetail = false)
         {
             return await InvokeAPIGatewayRequestWithContent(new TestLambdaContext(), GetRequestContent(fileName), configureApiToReturnExceptionDetail);
@@ -310,4 +336,19 @@ private string GetRequestContent(string fileName)
             return requestStr;
         }
     }
+
+    public class EnvironmentVariableHelper : IDisposable
+    {
+        private string _name;
+        private string? _oldValue;
+        public EnvironmentVariableHelper(string name, string value)
+        {
+            _name = name;
+            _oldValue = Environment.GetEnvironmentVariable(name);
+
+            Environment.SetEnvironmentVariable(name, value);
+        }
+
+        public void Dispose() => Environment.SetEnvironmentVariable(_name, _oldValue);
+    }
 }

Libraries/test/TestWebApp/Controllers/SnapStartController.cs

@@ -0,0 +1,24 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+using Microsoft.AspNetCore.Mvc;
+
+namespace TestWebApp.Controllers;
+
+[Route("api/[controller]")]
+public class SnapStartController
+{
+    /// <summary>
+    /// Set when <see cref="Get"/> is invoked
+    /// </summary>
+    public static bool Invoked { get; set; }
+
+
+    [HttpGet]
+    public string Get()
+    {
+        Invoked = true;
+
+        return "Invoked set to true";
+    }
+}

Libraries/test/TestWebApp/HttpApiV2LambdaFunction.cs

@@ -1,7 +1,16 @@
+using System.Collections.Generic;
+using System.Net.Http;
+using System;
 using Amazon.Lambda.AspNetCoreServer;
 namespace TestWebApp
 {
     public class HttpV2LambdaFunction : APIGatewayHttpApiV2ProxyFunction<Startup>
     {
+#if NET8_0_OR_GREATER
+        protected override IEnumerable<HttpRequestMessage> GetBeforeSnapshotRequests() =>
+        [
+            new HttpRequestMessage(HttpMethod.Get, "api/SnapStart")
+        ];
+#endif
     }
 }

Libraries/test/TestWebApp/Middleware.cs

@@ -1,4 +1,4 @@
-using Amazon.Lambda.Core;
+using Amazon.Lambda.Core;
 using Microsoft.AspNetCore.Http;
 using System;
 using System.Collections.Generic;
@@ -23,8 +23,12 @@ public async Task Invoke(HttpContext context)
 
             context.Response.OnStarting(x =>
             {
-                var lambdaContext = context.Items["LambdaContext"] as ILambdaContext;
-                lambdaContext?.Logger.LogLine("OnStarting Called");
+                if (context.Items.ContainsKey("LambdaContext"))
+                {
+                    var lambdaContext = context.Items["LambdaContext"] as ILambdaContext;
+                    lambdaContext?.Logger.LogLine("OnStarting Called");
+                }
+
                 return Task.FromResult(0);
             }, context);