Work on tests

Author: normj

Libraries/test/Amazon.Lambda.RuntimeSupport.Tests/Amazon.Lambda.RuntimeSupport.UnitTests/StreamingE2EWithMoq.cs

@@ -262,31 +262,42 @@ public async Task Buffered_ResponseBodyTransmittedCorrectly()
         [Fact]
         public async Task MidstreamError_SetsErrorStateWithExceptionDetails()
         {
-            var client = CreateClient();
-            const string errorMessage = "something went wrong mid-stream";
-
-            LambdaBootstrapHandler handler = async (invocation) =>
+            var testSuccess = false;
+            for (int i = 0; i < 5 && !testSuccess; i++)
             {
-                var stream = ResponseStreamFactory.CreateStream(Array.Empty<byte>());
-                await stream.WriteAsync(Encoding.UTF8.GetBytes("some data"));
-                await Task.Delay(1000);
-                throw new InvalidOperationException(errorMessage);
-            };
+                ResponseStreamFactory.CleanupInvocation(isMultiConcurrency: false);
+                var client = CreateClient();
+                const string errorMessage = "something went wrong mid-stream";
 
-            using var bootstrap = new LambdaBootstrap(handler, null);
-            bootstrap.Client = client;
-            await bootstrap.InvokeOnceAsync();
+                LambdaBootstrapHandler handler = async (invocation) =>
+                {
+                    var stream = ResponseStreamFactory.CreateStream(Array.Empty<byte>());
+                    await stream.WriteAsync(Encoding.UTF8.GetBytes("some data"));
+                    await Task.Delay(1000);
+                    throw new InvalidOperationException(errorMessage);
+                };
 
-            Assert.True(client.StartStreamingCalled);
-            Assert.NotNull(client.LastResponseStream);
-            Assert.True(client.LastResponseStream.HasError);
-            Assert.NotNull(client.LastResponseStream.ReportedError);
-            Assert.IsType<InvalidOperationException>(client.LastResponseStream.ReportedError);
-            Assert.Equal(errorMessage, client.LastResponseStream.ReportedError.Message);
-
-            // Verify the handler's data was still captured before the error
-            var output = Encoding.UTF8.GetString(client.CapturedHttpBytes);
-            Assert.Contains("some data", output);
+                using var bootstrap = new LambdaBootstrap(handler, null);
+                bootstrap.Client = client;
+                await bootstrap.InvokeOnceAsync();
+
+                if (!client.StartStreamingCalled)
+                    continue;
+                
+                Assert.NotNull(client.LastResponseStream);
+                Assert.True(client.LastResponseStream.HasError);
+                Assert.NotNull(client.LastResponseStream.ReportedError);
+                Assert.IsType<InvalidOperationException>(client.LastResponseStream.ReportedError);
+                Assert.Equal(errorMessage, client.LastResponseStream.ReportedError.Message);
+
+                // Verify the handler's data was still captured before the error
+                var output = Encoding.UTF8.GetString(client.CapturedHttpBytes);
+                Assert.Contains("some data", output);
+
+                testSuccess = true;
+            }
+
+            Assert.True(testSuccess);
         }
 
         // ─── 10.4 Multi-concurrency ──────────────────────────────────────────────────