Allow server longer to start

Give the server as long as it needs to start, then cancel after started.
martincostello · May 19, 2024 · 4647513 · 4647513
1 parent abf4ba8
commit 4647513
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 5 deletions.
diff --git a/tests/AwsLambdaTestServer.Tests/AwsIntegrationTests.cs b/tests/AwsLambdaTestServer.Tests/AwsIntegrationTests.cs
@@ -19,10 +19,12 @@ public static async Task Runtime_Generates_Valid_Aws_Trace_Id()
         Xunit.Skip.If(GetAwsCredentials() is null, "No AWS credentials are configured.");
 
         using var server = new LambdaTestServer();
-        using var cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+        using var cancellationTokenSource = new CancellationTokenSource();
 
         await server.StartAsync(cancellationTokenSource.Token);
 
+        cancellationTokenSource.CancelAfter(TimeSpan.FromSeconds(5));
+
         var request = new QueueExistsRequest()
         {
             QueueName = Guid.NewGuid().ToString(),

diff --git a/tests/AwsLambdaTestServer.Tests/Examples.cs b/tests/AwsLambdaTestServer.Tests/Examples.cs
@@ -18,11 +18,14 @@ public static async Task Function_Can_Process_Request()
 
         // Create a cancellation token that stops the server listening for new requests.
         // Auto-cancel the server after 2 seconds in case something goes wrong and the request is not handled.
-        using var cancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(2));
+        using var cancellationTokenSource = new CancellationTokenSource();
 
         // Start the test server so it is ready to listen for requests from the Lambda runtime
         await server.StartAsync(cancellationTokenSource.Token);
 
+        // Now that the server has started, cancel it after 2 seconds if no requests are processed
+        cancellationTokenSource.CancelAfter(TimeSpan.FromSeconds(2));
+
         // Create a test request for the Lambda function being tested
         var value = new MyRequest()
         {

diff --git a/tests/AwsLambdaTestServer.Tests/HttpLambdaTestServerTests.cs b/tests/AwsLambdaTestServer.Tests/HttpLambdaTestServerTests.cs
@@ -23,10 +23,12 @@ void Configure(IServiceCollection services)
             => services.AddLogging((builder) => builder.AddXUnit(this));
 
         using var server = new HttpLambdaTestServer(Configure);
-        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(2));
+        using var cts = new CancellationTokenSource();
 
         await server.StartAsync(cts.Token);
 
+        cts.CancelAfter(TimeSpan.FromSeconds(2));
+
         var context = await server.EnqueueAsync(@"{""Values"": [ 1, 2, 3 ]}");
 
         _ = Task.Run(async () =>
@@ -62,10 +64,12 @@ void Configure(IServiceCollection services)
             => services.AddLogging((builder) => builder.AddXUnit(this));
 
         using var server = new LambdaTestServer(Configure);
-        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(2));
+        using var cts = new CancellationTokenSource();
 
         await server.StartAsync(cts.Token);
 
+        cts.CancelAfter(TimeSpan.FromSeconds(2));
+
         var context = await server.EnqueueAsync(@"{""Values"": null}");
 
         _ = Task.Run(async () =>
@@ -99,10 +103,12 @@ void Configure(IServiceCollection services)
             => services.AddLogging((builder) => builder.AddXUnit(this));
 
         using var server = new LambdaTestServer(Configure);
-        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(2));
+        using var cts = new CancellationTokenSource();
 
         await server.StartAsync(cts.Token);
 
+        cts.CancelAfter(TimeSpan.FromSeconds(2));
+
         var channels = new List<(int Expected, LambdaTestContext Context)>();
 
         for (int i = 0; i < 10; i++)