chore: improve test flakiness

philasmar · philasmar · commit 900e67bc2efd · 2026-04-08T20:29:04.000-04:00
diff --git a/.autover/changes/9cf844c3-1a90-46ec-9430-93724dcb4512.json b/.autover/changes/9cf844c3-1a90-46ec-9430-93724dcb4512.json
@@ -0,0 +1,11 @@
+{
+  "Projects": [
+    {
+      "Name": "Amazon.Lambda.TestTool.BlazorTester",
+      "Type": "Patch",
+      "ChangelogMessages": [
+        "Minor fixes and improvements"
+      ]
+    }
+  ]
+}
diff --git a/.autover/changes/a119819e-ba57-424d-a0be-d941eea599f9.json b/.autover/changes/a119819e-ba57-424d-a0be-d941eea599f9.json
@@ -0,0 +1,11 @@
+{
+  "Projects": [
+    {
+      "Name": "Amazon.Lambda.RuntimeSupport",
+      "Type": "Patch",
+      "ChangelogMessages": [
+        "Minor fixes and improvements"
+      ]
+    }
+  ]
+}
diff --git a/Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/ConsoleLoggerWriter.cs b/Libraries/src/Amazon.Lambda.RuntimeSupport/Helpers/ConsoleLoggerWriter.cs
@@ -227,6 +227,7 @@ public LogLevelLoggerWriter(IEnvironmentVariables environmentVariables)
         /// <param name="stdErrorWriter"></param>
         public LogLevelLoggerWriter(TextWriter stdOutWriter, TextWriter stdErrorWriter)
         {
+            _environmentVariables = new SystemEnvironmentVariables();
             Initialize(stdOutWriter, stdErrorWriter);
         }
 
@@ -325,15 +326,15 @@ public IRuntimeApiHeaders CurrentRuntimeApiHeaders
             {
                 get
                 {
-                    if (Utils.IsUsingMultiConcurrency(_environmentVariables))
+                    if (_currentRuntimeApiHeadersStorage != null && Utils.IsUsingMultiConcurrency(_environmentVariables))
                     {
                         return _currentRuntimeApiHeadersStorage.Value;
                     }
                     return _currentRuntimeApiHeaders;
                 }
                 set
                 {
-                    if (Utils.IsUsingMultiConcurrency(_environmentVariables))
+                    if (_currentRuntimeApiHeadersStorage != null && Utils.IsUsingMultiConcurrency(_environmentVariables))
                     {
                         _currentRuntimeApiHeadersStorage.Value = value;
                     }
diff --git a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/HandlerTests.cs b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/HandlerTests.cs
@@ -250,7 +250,7 @@ private async Task<Exception> TestHandlerFailAsync(string handler, string expect
             var userCodeLoader = new UserCodeLoader(new SystemEnvironmentVariables(), handler, _internalLogger);
             var initializer = new UserCodeInitializer(userCodeLoader, _internalLogger);
             var handlerWrapper = HandlerWrapper.GetHandlerWrapper(userCodeLoader.Invoke);
-            var bootstrap = new LambdaBootstrap(handlerWrapper, initializer.InitializeAsync)
+            var bootstrap = new LambdaBootstrap(handlerWrapper.Handler, initializer.InitializeAsync, null, _environmentVariables)
             {
                 Client = testRuntimeApiClient
             };
@@ -388,7 +388,9 @@ private async Task<ExecutionInfo> ExecHandlerAsync(string handler, string dataIn
                 var userCodeLoader = new UserCodeLoader(new SystemEnvironmentVariables(), handler, _internalLogger);
                 var handlerWrapper = HandlerWrapper.GetHandlerWrapper(userCodeLoader.Invoke);
                 var initializer = new UserCodeInitializer(userCodeLoader, _internalLogger);
-                var bootstrap = new LambdaBootstrap(handlerWrapper, initializer.InitializeAsync)
+                // Pass null initializer to bootstrap so RunAsync won't re-invoke Init(),
+                // which would re-register AssemblyLoad event handlers and re-construct the invoke delegate.
+                var bootstrap = new LambdaBootstrap(handlerWrapper.Handler, null, null, _environmentVariables)
                 {
                     Client = testRuntimeApiClient
                 };
@@ -403,7 +405,13 @@ private async Task<ExecutionInfo> ExecHandlerAsync(string handler, string dataIn
                     Assert.DoesNotContain($"^^[{assertLoggedByInitialize}]^^", actionWriter.ToString());
                 }
 
-                await bootstrap.InitializeAsync();
+                await initializer.InitializeAsync();
+
+                // Re-set logging actions after initialization in case Init's AssemblyLoad event
+                // handler overwrote them when loading Amazon.Lambda.Core as a handler dependency.
+                UserCodeLoader.SetCustomerLoggerLogAction(assembly, actionWriter.ToLoggingAction(), _internalLogger);
+                UserCodeLoader.SetCustomerLoggerLogAction(assembly, actionWriter.ToLoggingWithLevelAction(), _internalLogger);
+                UserCodeLoader.SetCustomerLoggerLogAction(assembly, actionWriter.ToLoggingWithLevelAndExceptionAction(), _internalLogger);
 
                 if (assertLoggedByInitialize != null)
                 {
diff --git a/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/LambdaBootstrapTests.cs b/Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/LambdaBootstrapTests.cs
@@ -165,7 +165,7 @@ public async Task TraceIdEnvironmentVariableIsSet()
         [Fact]
         public async Task HandlerThrowsException()
         {
-            using (var bootstrap = new LambdaBootstrap(_testFunction.BaseHandlerThrowsAsync, null))
+            using (var bootstrap = new LambdaBootstrap(_testFunction.BaseHandlerThrowsAsync, null, null, _environmentVariables))
             {
                 bootstrap.Client = _testRuntimeApiClient;
                 Assert.Null(_environmentVariables.GetEnvironmentVariable(LambdaEnvironment.EnvVarTraceId));
@@ -183,7 +183,7 @@ public async Task HandlerInputAndOutputWork()
         {
             const string testInput = "a MiXeD cAsE sTrInG";
 
-            using (var bootstrap = new LambdaBootstrap(_testFunction.BaseHandlerToUpperAsync, null))
+            using (var bootstrap = new LambdaBootstrap(_testFunction.BaseHandlerToUpperAsync, null, null, _environmentVariables))
             {
                 _testRuntimeApiClient.FunctionInput = Encoding.UTF8.GetBytes(testInput);
                 bootstrap.Client = _testRuntimeApiClient;
@@ -201,7 +201,7 @@ public async Task HandlerInputAndOutputWork()
         [Fact]
         public async Task HandlerReturnsNull()
         {
-            using (var bootstrap = new LambdaBootstrap(_testFunction.BaseHandlerReturnsNullAsync, null))
+            using (var bootstrap = new LambdaBootstrap(_testFunction.BaseHandlerReturnsNullAsync, null, null, _environmentVariables))
             {
                 _testRuntimeApiClient.FunctionInput = new byte[0];
                 bootstrap.Client = _testRuntimeApiClient;
diff --git a/Tools/LambdaTestTool-v2/tests/Amazon.Lambda.TestTool.IntegrationTests/BaseApiGatewayTest.cs b/Tools/LambdaTestTool-v2/tests/Amazon.Lambda.TestTool.IntegrationTests/BaseApiGatewayTest.cs
@@ -48,6 +48,7 @@ protected async Task CleanupAsync()
             CancellationTokenSource.Dispose();
             CancellationTokenSource = new CancellationTokenSource();
         }
+        Environment.SetEnvironmentVariable("APIGATEWAY_EMULATOR_ROUTE_CONFIG", null);
     }
 
     protected async Task StartTestToolProcessAsync(ApiGatewayEmulatorMode apiGatewayMode, string routeName, int lambdaPort, int apiGatewayPort, CancellationTokenSource cancellationTokenSource, string httpMethod = "POST")
diff --git a/Tools/LambdaTestTool-v2/tests/Amazon.Lambda.TestTool.IntegrationTests/SQSEventSourceTests.cs b/Tools/LambdaTestTool-v2/tests/Amazon.Lambda.TestTool.IntegrationTests/SQSEventSourceTests.cs
@@ -183,7 +183,7 @@ public async Task ProcessMessagesFromMultipleEventSources()
             await sqsClient.SendMessageAsync(queueUrl2, "MessageFromQueue2");
 
             var startTime = DateTime.UtcNow;
-            while (listOfProcessedMessages.Count == 0 && DateTime.UtcNow < startTime.AddMinutes(2))
+            while (listOfProcessedMessages.Count < 2 && DateTime.UtcNow < startTime.AddMinutes(2))
             {
                 Assert.False(lambdaTask.IsFaulted, "Lambda function failed: " + lambdaTask.Exception?.ToString());
                 await Task.Delay(500);
diff --git a/Tools/LambdaTestTool-v2/tests/Amazon.Lambda.TestTool.Tests.Common/Helpers/TestHelpers.cs b/Tools/LambdaTestTool-v2/tests/Amazon.Lambda.TestTool.Tests.Common/Helpers/TestHelpers.cs
@@ -1,6 +1,9 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+using System.Net;
+using System.Net.Sockets;
+
 namespace Amazon.Lambda.TestTool.Tests.Common.Helpers;
 
 public static class TestHelpers
@@ -39,16 +42,22 @@ public static async Task<HttpResponseMessage> SendRequest(string url)
         }
     }
 
-    private static int _maxLambdaRuntimePort = 6000;
-    private static int _maxApiGatewayPort = 9000;
+    private static int GetFreePort()
+    {
+        var listener = new TcpListener(IPAddress.Loopback, 0);
+        listener.Start();
+        var port = ((IPEndPoint)listener.LocalEndpoint).Port;
+        listener.Stop();
+        return port;
+    }
 
     public static int GetNextLambdaRuntimePort()
     {
-        return Interlocked.Increment(ref _maxLambdaRuntimePort);
+        return GetFreePort();
     }
 
     public static int GetNextApiGatewayPort()
     {
-        return Interlocked.Increment(ref _maxApiGatewayPort);
+        return GetFreePort();
     }
 }
diff --git a/Tools/LambdaTestTool-v2/tests/Amazon.Lambda.TestTool.UnitTests/RuntimeApiTests.cs b/Tools/LambdaTestTool-v2/tests/Amazon.Lambda.TestTool.UnitTests/RuntimeApiTests.cs
@@ -39,6 +39,8 @@ public async Task AddEventToDataStore()
         var testToolProcess = TestToolProcess.Startup(options, cancellationTokenSource.Token);
         try
         {
+            Assert.True(await TestHelpers.WaitForApiToStartAsync($"{testToolProcess.ServiceUrl}/lambda-runtime-api/healthcheck"));
+
             var lambdaClient = ConstructLambdaServiceClient(testToolProcess.ServiceUrl);
             var invokeFunction = new InvokeRequest
             {
@@ -92,6 +94,8 @@ public async Task InvokeRequestResponse()
         var testToolProcess = TestToolProcess.Startup(options, cancellationTokenSource.Token);
         try
         {
+            Assert.True(await TestHelpers.WaitForApiToStartAsync($"{testToolProcess.ServiceUrl}/lambda-runtime-api/healthcheck"));
+
             var handler = (string input, ILambdaContext context) =>
             {
                 Thread.Sleep(1000); // Add a sleep to prove the LambdaRuntimeApi waited for the completion.
diff --git a/Tools/LambdaTestTool/src/Amazon.Lambda.TestTool/TestToolStartup.cs b/Tools/LambdaTestTool/src/Amazon.Lambda.TestTool/TestToolStartup.cs
@@ -10,7 +10,6 @@ namespace Amazon.Lambda.TestTool
 {
     public class TestToolStartup
     {
-        private static bool shouldDisableLogs;
 
         public class RunConfiguration
         {
@@ -37,7 +36,7 @@ public static void Startup(string productName, Action<LocalLambdaOptions, bool>
             try
             {
                 var commandOptions = CommandLineOptions.Parse(args);
-                shouldDisableLogs = Utils.ShouldDisableLogs(commandOptions);
+                var shouldDisableLogs = Utils.ShouldDisableLogs(commandOptions);
 
                 if (!shouldDisableLogs) Utils.PrintToolTitle(productName);
 
@@ -76,7 +75,7 @@ public static void Startup(string productName, Action<LocalLambdaOptions, bool>
 
                 if (commandOptions.NoUI)
                 {
-                    ExecuteWithNoUi(localLambdaOptions, commandOptions, lambdaAssemblyDirectory, runConfiguration);
+                    ExecuteWithNoUi(localLambdaOptions, commandOptions, lambdaAssemblyDirectory, runConfiguration, shouldDisableLogs);
                 }
                 else
                 {
@@ -118,16 +117,16 @@ public static void Startup(string productName, Action<LocalLambdaOptions, bool>
         }
 
 
-        public static void ExecuteWithNoUi(LocalLambdaOptions localLambdaOptions, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, RunConfiguration runConfiguration)
+        public static void ExecuteWithNoUi(LocalLambdaOptions localLambdaOptions, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, RunConfiguration runConfiguration, bool shouldDisableLogs)
         {
             if (!shouldDisableLogs) runConfiguration.OutputWriter.WriteLine("Executing Lambda function without web interface");
             var lambdaProjectDirectory = Utils.FindLambdaProjectDirectory(lambdaAssemblyDirectory);
 
-            string configFile = DetermineConfigFile(commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory);
-            LambdaConfigInfo configInfo = LoadLambdaConfigInfo(configFile, commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory, runConfiguration);
-            LambdaFunction lambdaFunction = LoadLambdaFunction(configInfo, localLambdaOptions, commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory, runConfiguration);
+            string configFile = DetermineConfigFile(commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory, shouldDisableLogs: shouldDisableLogs);
+            LambdaConfigInfo configInfo = LoadLambdaConfigInfo(configFile, commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory, runConfiguration, shouldDisableLogs: shouldDisableLogs);
+            LambdaFunction lambdaFunction = LoadLambdaFunction(configInfo, localLambdaOptions, commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory, runConfiguration, shouldDisableLogs: shouldDisableLogs);
 
-            string payload = DeterminePayload(localLambdaOptions, commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory, runConfiguration);
+            string payload = DeterminePayload(localLambdaOptions, commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory, runConfiguration, shouldDisableLogs: shouldDisableLogs);
 
             var awsProfile = commandOptions.AWSProfile ?? configInfo.AWSProfile;
             if (!string.IsNullOrEmpty(awsProfile))
@@ -166,7 +165,7 @@ public static void ExecuteWithNoUi(LocalLambdaOptions localLambdaOptions, Comman
                 Function = lambdaFunction
             };
 
-            ExecuteRequest(request, localLambdaOptions, runConfiguration);
+            ExecuteRequest(request, localLambdaOptions, runConfiguration, shouldDisableLogs);
 
 
             if (runConfiguration.Mode == RunConfiguration.RunMode.Normal && commandOptions.PauseExit)
@@ -176,7 +175,7 @@ public static void ExecuteWithNoUi(LocalLambdaOptions localLambdaOptions, Comman
             }
         }
 
-        private static string DetermineConfigFile(CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory)
+        private static string DetermineConfigFile(CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory, bool shouldDisableLogs)
         {
             string configFile = null;
             if (string.IsNullOrEmpty(commandOptions.ConfigFile))
@@ -199,7 +198,7 @@ private static string DetermineConfigFile(CommandLineOptions commandOptions, str
             return configFile;
         }
 
-        private static LambdaConfigInfo LoadLambdaConfigInfo(string configFile, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory, RunConfiguration runConfiguration)
+        private static LambdaConfigInfo LoadLambdaConfigInfo(string configFile, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory, RunConfiguration runConfiguration, bool shouldDisableLogs)
         {
             LambdaConfigInfo configInfo;
             if (configFile != null)
@@ -226,7 +225,7 @@ private static LambdaConfigInfo LoadLambdaConfigInfo(string configFile, CommandL
             return configInfo;
         }
 
-        private static LambdaFunction LoadLambdaFunction(LambdaConfigInfo configInfo, LocalLambdaOptions localLambdaOptions, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory, RunConfiguration runConfiguration)
+        private static LambdaFunction LoadLambdaFunction(LambdaConfigInfo configInfo, LocalLambdaOptions localLambdaOptions, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory, RunConfiguration runConfiguration, bool shouldDisableLogs)
         {
             // If no function handler was explicitly set and there is only one function defined in the config file then assume the user wants to debug that function.
             var functionHandler = commandOptions.FunctionHandler;
@@ -264,7 +263,7 @@ private static LambdaFunction LoadLambdaFunction(LambdaConfigInfo configInfo, Lo
             return lambdaFunction;
         }
 
-        private static string DeterminePayload(LocalLambdaOptions localLambdaOptions, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory, RunConfiguration runConfiguration)
+        private static string DeterminePayload(LocalLambdaOptions localLambdaOptions, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory, RunConfiguration runConfiguration, bool shouldDisableLogs)
         {
             var payload = commandOptions.Payload;
 
@@ -346,7 +345,7 @@ private static string DeterminePayload(LocalLambdaOptions localLambdaOptions, Co
             return payload;
         }
 
-        private static void ExecuteRequest(ExecutionRequest request, LocalLambdaOptions localLambdaOptions, RunConfiguration runConfiguration)
+        private static void ExecuteRequest(ExecutionRequest request, LocalLambdaOptions localLambdaOptions, RunConfiguration runConfiguration, bool shouldDisableLogs)
         {
             try
             {

Original file line number	Diff line number	Diff line change
`@@ -227,6 +227,7 @@ public LogLevelLoggerWriter(IEnvironmentVariables environmentVariables)`
`227`	`227`	`/// <param name="stdErrorWriter"></param>`
`228`	`228`	`public LogLevelLoggerWriter(TextWriter stdOutWriter, TextWriter stdErrorWriter)`
`229`	`229`	`{`
	`230`	`+ _environmentVariables = new SystemEnvironmentVariables();`
`230`	`231`	`Initialize(stdOutWriter, stdErrorWriter);`
`231`	`232`	`}`
`232`	`233`
`@@ -325,15 +326,15 @@ public IRuntimeApiHeaders CurrentRuntimeApiHeaders`
`325`	`326`	`{`
`326`	`327`	`get`
`327`	`328`	`{`
`328`		`- if (Utils.IsUsingMultiConcurrency(_environmentVariables))`
	`329`	`+ if (_currentRuntimeApiHeadersStorage != null && Utils.IsUsingMultiConcurrency(_environmentVariables))`
`329`	`330`	`{`
`330`	`331`	`return _currentRuntimeApiHeadersStorage.Value;`
`331`	`332`	`}`
`332`	`333`	`return _currentRuntimeApiHeaders;`
`333`	`334`	`}`
`334`	`335`	`set`
`335`	`336`	`{`
`336`		`- if (Utils.IsUsingMultiConcurrency(_environmentVariables))`
	`337`	`+ if (_currentRuntimeApiHeadersStorage != null && Utils.IsUsingMultiConcurrency(_environmentVariables))`
`337`	`338`	`{`
`338`	`339`	`_currentRuntimeApiHeadersStorage.Value = value;`
`339`	`340`	`}`
Original file line number	Diff line number	Diff line change
`@@ -165,7 +165,7 @@ public async Task TraceIdEnvironmentVariableIsSet()`
`165`	`165`	`[Fact]`
`166`	`166`	`public async Task HandlerThrowsException()`
`167`	`167`	`{`
`168`		`- using (var bootstrap = new LambdaBootstrap(_testFunction.BaseHandlerThrowsAsync, null))`
	`168`	`+ using (var bootstrap = new LambdaBootstrap(_testFunction.BaseHandlerThrowsAsync, null, null, _environmentVariables))`
`169`	`169`	`{`
`170`	`170`	`bootstrap.Client = _testRuntimeApiClient;`
`171`	`171`	`Assert.Null(_environmentVariables.GetEnvironmentVariable(LambdaEnvironment.EnvVarTraceId));`
`@@ -183,7 +183,7 @@ public async Task HandlerInputAndOutputWork()`
`183`	`183`	`{`
`184`	`184`	`const string testInput = "a MiXeD cAsE sTrInG";`
`185`	`185`
`186`		`- using (var bootstrap = new LambdaBootstrap(_testFunction.BaseHandlerToUpperAsync, null))`
	`186`	`+ using (var bootstrap = new LambdaBootstrap(_testFunction.BaseHandlerToUpperAsync, null, null, _environmentVariables))`
`187`	`187`	`{`
`188`	`188`	`_testRuntimeApiClient.FunctionInput = Encoding.UTF8.GetBytes(testInput);`
`189`	`189`	`bootstrap.Client = _testRuntimeApiClient;`
`@@ -201,7 +201,7 @@ public async Task HandlerInputAndOutputWork()`
`201`	`201`	`[Fact]`
`202`	`202`	`public async Task HandlerReturnsNull()`
`203`	`203`	`{`
`204`		`- using (var bootstrap = new LambdaBootstrap(_testFunction.BaseHandlerReturnsNullAsync, null))`
	`204`	`+ using (var bootstrap = new LambdaBootstrap(_testFunction.BaseHandlerReturnsNullAsync, null, null, _environmentVariables))`
`205`	`205`	`{`
`206`	`206`	`_testRuntimeApiClient.FunctionInput = new byte[0];`
`207`	`207`	`bootstrap.Client = _testRuntimeApiClient;`
Original file line number	Diff line number	Diff line change
`@@ -48,6 +48,7 @@ protected async Task CleanupAsync()`
`48`	`48`	`CancellationTokenSource.Dispose();`
`49`	`49`	`CancellationTokenSource = new CancellationTokenSource();`
`50`	`50`	`}`
	`51`	`+ Environment.SetEnvironmentVariable("APIGATEWAY_EMULATOR_ROUTE_CONFIG", null);`
`51`	`52`	`}`
`52`	`53`
`53`	`54`	`protected async Task StartTestToolProcessAsync(ApiGatewayEmulatorMode apiGatewayMode, string routeName, int lambdaPort, int apiGatewayPort, CancellationTokenSource cancellationTokenSource, string httpMethod = "POST")`
Original file line number	Diff line number	Diff line change
`@@ -183,7 +183,7 @@ public async Task ProcessMessagesFromMultipleEventSources()`
`183`	`183`	`await sqsClient.SendMessageAsync(queueUrl2, "MessageFromQueue2");`
`184`	`184`
`185`	`185`	`var startTime = DateTime.UtcNow;`
`186`		`- while (listOfProcessedMessages.Count == 0 && DateTime.UtcNow < startTime.AddMinutes(2))`
	`186`	`+ while (listOfProcessedMessages.Count < 2 && DateTime.UtcNow < startTime.AddMinutes(2))`
`187`	`187`	`{`
`188`	`188`	`Assert.False(lambdaTask.IsFaulted, "Lambda function failed: " + lambdaTask.Exception?.ToString());`
`189`	`189`	`await Task.Delay(500);`
Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,9 @@`
`1`	`1`	`// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.`
`2`	`2`	`// SPDX-License-Identifier: Apache-2.0`
`3`	`3`
	`4`	`+using System.Net;`
	`5`	`+using System.Net.Sockets;`
	`6`	`+`
`4`	`7`	`namespace Amazon.Lambda.TestTool.Tests.Common.Helpers;`
`5`	`8`
`6`	`9`	`public static class TestHelpers`
`@@ -39,16 +42,22 @@ public static async Task<HttpResponseMessage> SendRequest(string url)`
`39`	`42`	`}`
`40`	`43`	`}`
`41`	`44`
`42`		`- private static int _maxLambdaRuntimePort = 6000;`
`43`		`- private static int _maxApiGatewayPort = 9000;`
	`45`	`+ private static int GetFreePort()`
	`46`	`+ {`
	`47`	`+ var listener = new TcpListener(IPAddress.Loopback, 0);`
	`48`	`+ listener.Start();`
	`49`	`+ var port = ((IPEndPoint)listener.LocalEndpoint).Port;`
	`50`	`+ listener.Stop();`
	`51`	`+ return port;`
	`52`	`+ }`
`44`	`53`
`45`	`54`	`public static int GetNextLambdaRuntimePort()`
`46`	`55`	`{`
`47`		`- return Interlocked.Increment(ref _maxLambdaRuntimePort);`
	`56`	`+ return GetFreePort();`
`48`	`57`	`}`
`49`	`58`
`50`	`59`	`public static int GetNextApiGatewayPort()`
`51`	`60`	`{`
`52`		`- return Interlocked.Increment(ref _maxApiGatewayPort);`
	`61`	`+ return GetFreePort();`
`53`	`62`	`}`
`54`	`63`	`}`
Original file line number	Diff line number	Diff line change
`@@ -39,6 +39,8 @@ public async Task AddEventToDataStore()`
`39`	`39`	`var testToolProcess = TestToolProcess.Startup(options, cancellationTokenSource.Token);`
`40`	`40`	`try`
`41`	`41`	`{`
	`42`	`+ Assert.True(await TestHelpers.WaitForApiToStartAsync($"{testToolProcess.ServiceUrl}/lambda-runtime-api/healthcheck"));`
	`43`	`+`
`42`	`44`	`var lambdaClient = ConstructLambdaServiceClient(testToolProcess.ServiceUrl);`
`43`	`45`	`var invokeFunction = new InvokeRequest`
`44`	`46`	`{`
`@@ -92,6 +94,8 @@ public async Task InvokeRequestResponse()`
`92`	`94`	`var testToolProcess = TestToolProcess.Startup(options, cancellationTokenSource.Token);`
`93`	`95`	`try`
`94`	`96`	`{`
	`97`	`+ Assert.True(await TestHelpers.WaitForApiToStartAsync($"{testToolProcess.ServiceUrl}/lambda-runtime-api/healthcheck"));`
	`98`	`+`
`95`	`99`	`var handler = (string input, ILambdaContext context) =>`
`96`	`100`	`{`
`97`	`101`	`Thread.Sleep(1000); // Add a sleep to prove the LambdaRuntimeApi waited for the completion.`
Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,6 @@ namespace Amazon.Lambda.TestTool`
`10`	`10`	`{`
`11`	`11`	`public class TestToolStartup`
`12`	`12`	`{`
`13`		`- private static bool shouldDisableLogs;`
`14`	`13`
`15`	`14`	`public class RunConfiguration`
`16`	`15`	`{`
`@@ -37,7 +36,7 @@ public static void Startup(string productName, Action<LocalLambdaOptions, bool>`
`37`	`36`	`try`
`38`	`37`	`{`
`39`	`38`	`var commandOptions = CommandLineOptions.Parse(args);`
`40`		`- shouldDisableLogs = Utils.ShouldDisableLogs(commandOptions);`
	`39`	`+ var shouldDisableLogs = Utils.ShouldDisableLogs(commandOptions);`
`41`	`40`
`42`	`41`	`if (!shouldDisableLogs) Utils.PrintToolTitle(productName);`
`43`	`42`
`@@ -76,7 +75,7 @@ public static void Startup(string productName, Action<LocalLambdaOptions, bool>`
`76`	`75`
`77`	`76`	`if (commandOptions.NoUI)`
`78`	`77`	`{`
`79`		`- ExecuteWithNoUi(localLambdaOptions, commandOptions, lambdaAssemblyDirectory, runConfiguration);`
	`78`	`+ ExecuteWithNoUi(localLambdaOptions, commandOptions, lambdaAssemblyDirectory, runConfiguration, shouldDisableLogs);`
`80`	`79`	`}`
`81`	`80`	`else`
`82`	`81`	`{`
`@@ -118,16 +117,16 @@ public static void Startup(string productName, Action<LocalLambdaOptions, bool>`
`118`	`117`	`}`
`119`	`118`
`120`	`119`
`121`		`- public static void ExecuteWithNoUi(LocalLambdaOptions localLambdaOptions, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, RunConfiguration runConfiguration)`
	`120`	`+ public static void ExecuteWithNoUi(LocalLambdaOptions localLambdaOptions, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, RunConfiguration runConfiguration, bool shouldDisableLogs)`
`122`	`121`	`{`
`123`	`122`	`if (!shouldDisableLogs) runConfiguration.OutputWriter.WriteLine("Executing Lambda function without web interface");`
`124`	`123`	`var lambdaProjectDirectory = Utils.FindLambdaProjectDirectory(lambdaAssemblyDirectory);`
`125`	`124`
`126`		`- string configFile = DetermineConfigFile(commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory);`
`127`		`- LambdaConfigInfo configInfo = LoadLambdaConfigInfo(configFile, commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory, runConfiguration);`
`128`		`- LambdaFunction lambdaFunction = LoadLambdaFunction(configInfo, localLambdaOptions, commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory, runConfiguration);`
	`125`	`+ string configFile = DetermineConfigFile(commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory, shouldDisableLogs: shouldDisableLogs);`
	`126`	`+ LambdaConfigInfo configInfo = LoadLambdaConfigInfo(configFile, commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory, runConfiguration, shouldDisableLogs: shouldDisableLogs);`
	`127`	`+ LambdaFunction lambdaFunction = LoadLambdaFunction(configInfo, localLambdaOptions, commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory, runConfiguration, shouldDisableLogs: shouldDisableLogs);`
`129`	`128`
`130`		`- string payload = DeterminePayload(localLambdaOptions, commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory, runConfiguration);`
	`129`	`+ string payload = DeterminePayload(localLambdaOptions, commandOptions, lambdaAssemblyDirectory: lambdaAssemblyDirectory, lambdaProjectDirectory: lambdaProjectDirectory, runConfiguration, shouldDisableLogs: shouldDisableLogs);`
`131`	`130`
`132`	`131`	`var awsProfile = commandOptions.AWSProfile ?? configInfo.AWSProfile;`
`133`	`132`	`if (!string.IsNullOrEmpty(awsProfile))`
`@@ -166,7 +165,7 @@ public static void ExecuteWithNoUi(LocalLambdaOptions localLambdaOptions, Comman`
`166`	`165`	`Function = lambdaFunction`
`167`	`166`	`};`
`168`	`167`
`169`		`- ExecuteRequest(request, localLambdaOptions, runConfiguration);`
	`168`	`+ ExecuteRequest(request, localLambdaOptions, runConfiguration, shouldDisableLogs);`
`170`	`169`
`171`	`170`
`172`	`171`	`if (runConfiguration.Mode == RunConfiguration.RunMode.Normal && commandOptions.PauseExit)`
`@@ -176,7 +175,7 @@ public static void ExecuteWithNoUi(LocalLambdaOptions localLambdaOptions, Comman`
`176`	`175`	`}`
`177`	`176`	`}`
`178`	`177`
`179`		`- private static string DetermineConfigFile(CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory)`
	`178`	`+ private static string DetermineConfigFile(CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory, bool shouldDisableLogs)`
`180`	`179`	`{`
`181`	`180`	`string configFile = null;`
`182`	`181`	`if (string.IsNullOrEmpty(commandOptions.ConfigFile))`
`@@ -199,7 +198,7 @@ private static string DetermineConfigFile(CommandLineOptions commandOptions, str`
`199`	`198`	`return configFile;`
`200`	`199`	`}`
`201`	`200`
`202`		`- private static LambdaConfigInfo LoadLambdaConfigInfo(string configFile, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory, RunConfiguration runConfiguration)`
	`201`	`+ private static LambdaConfigInfo LoadLambdaConfigInfo(string configFile, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory, RunConfiguration runConfiguration, bool shouldDisableLogs)`
`203`	`202`	`{`
`204`	`203`	`LambdaConfigInfo configInfo;`
`205`	`204`	`if (configFile != null)`
`@@ -226,7 +225,7 @@ private static LambdaConfigInfo LoadLambdaConfigInfo(string configFile, CommandL`
`226`	`225`	`return configInfo;`
`227`	`226`	`}`
`228`	`227`
`229`		`- private static LambdaFunction LoadLambdaFunction(LambdaConfigInfo configInfo, LocalLambdaOptions localLambdaOptions, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory, RunConfiguration runConfiguration)`
	`228`	`+ private static LambdaFunction LoadLambdaFunction(LambdaConfigInfo configInfo, LocalLambdaOptions localLambdaOptions, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory, RunConfiguration runConfiguration, bool shouldDisableLogs)`
`230`	`229`	`{`
`231`	`230`	`// If no function handler was explicitly set and there is only one function defined in the config file then assume the user wants to debug that function.`
`232`	`231`	`var functionHandler = commandOptions.FunctionHandler;`
`@@ -264,7 +263,7 @@ private static LambdaFunction LoadLambdaFunction(LambdaConfigInfo configInfo, Lo`
`264`	`263`	`return lambdaFunction;`
`265`	`264`	`}`
`266`	`265`
`267`		`- private static string DeterminePayload(LocalLambdaOptions localLambdaOptions, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory, RunConfiguration runConfiguration)`
	`266`	`+ private static string DeterminePayload(LocalLambdaOptions localLambdaOptions, CommandLineOptions commandOptions, string lambdaAssemblyDirectory, string lambdaProjectDirectory, RunConfiguration runConfiguration, bool shouldDisableLogs)`
`268`	`267`	`{`
`269`	`268`	`var payload = commandOptions.Payload;`
`270`	`269`
`@@ -346,7 +345,7 @@ private static string DeterminePayload(LocalLambdaOptions localLambdaOptions, Co`
`346`	`345`	`return payload;`
`347`	`346`	`}`
`348`	`347`
`349`		`- private static void ExecuteRequest(ExecutionRequest request, LocalLambdaOptions localLambdaOptions, RunConfiguration runConfiguration)`
	`348`	`+ private static void ExecuteRequest(ExecutionRequest request, LocalLambdaOptions localLambdaOptions, RunConfiguration runConfiguration, bool shouldDisableLogs)`
`350`	`349`	`{`
`351`	`350`	`try`
`352`	`351`	`{`