From 4758d4e6cc563a5f06181fe58bead1fb8e3cdb9d Mon Sep 17 00:00:00 2001
From: Alfred Neequaye <aneequaye@acsint.io>
Date: Fri, 2 Dec 2022 04:33:21 +0000
Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=EF=B8=8F=20Prevent=20Infinite=20Wait?=
 =?UTF-8?q?=20with=20Timeouts=20Exceptions=20and=20Retry=20counts=20in=20U?=
 =?UTF-8?q?ntilHttp=20Strategy?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../Configurations/UntilHttpOptions.cs        |  2 +
 .../WaitStrategies/UntilHttp.cs               | 68 ++++++++++++-------
 2 files changed, 45 insertions(+), 25 deletions(-)

diff --git a/src/Testcontainers/Configurations/UntilHttpOptions.cs b/src/Testcontainers/Configurations/UntilHttpOptions.cs
index d40083313..82687e8ee 100644
--- a/src/Testcontainers/Configurations/UntilHttpOptions.cs
+++ b/src/Testcontainers/Configurations/UntilHttpOptions.cs
@@ -21,6 +21,7 @@ public UntilHttpOptions()
       this.ExpectedResponseCodes = new() { HttpStatusCode.OK };
       this.TimeOut = TimeSpan.FromMinutes(1);
       this.RequestDelay = 1;
+      this.MaxRetries = 10;
     }
 
     public HttpMethod Method { get; set; }
@@ -36,6 +37,7 @@ public UntilHttpOptions()
     public TimeSpan TimeOut { get; set; }
     public bool ValidateContent { get; set; }
     public double RequestDelay { get; set; }
+    public int MaxRetries { get; set; }
 
     public Uri Uri => new($"{(this.UseSecure ? "https" : "http")}://{this.Host}:{this.Port}{this.Path}");
   }
diff --git a/src/Testcontainers/Configurations/WaitStrategies/UntilHttp.cs b/src/Testcontainers/Configurations/WaitStrategies/UntilHttp.cs
index 9ce907b66..ff8bfb933 100644
--- a/src/Testcontainers/Configurations/WaitStrategies/UntilHttp.cs
+++ b/src/Testcontainers/Configurations/WaitStrategies/UntilHttp.cs
@@ -12,6 +12,7 @@
   public class UntilHttp : IWaitUntil
   {
     private readonly UntilHttpOptions Options;
+    private int RetryCount = 0;
 
     public UntilHttp(string path)
     {
@@ -25,42 +26,59 @@ public UntilHttp(UntilHttpOptions options)
 
     public async Task<bool> Until(ITestcontainersContainer testcontainers, ILogger logger)
     {
-      var mappedPort = testcontainers.GetMappedPublicPort(this.Options.Port);
-      this.Options.Port = mappedPort;
-      var client = new HttpClient();
-      var message = new HttpRequestMessage(this.Options.Method, this.Options.Uri);
-      if (this.Options.RequestContent is not null && (this.Options.Method == HttpMethod.Post || this.Options.Method == HttpMethod.Put))
-      {
-        message.Content = this.Options.RequestContent;
-      }
 
-      if (this.Options.UseAuth && this.Options.AuthString is not null)
+      try
       {
-        message.Headers.Authorization = AuthenticationHeaderValue.Parse(this.Options.AuthString.ToString());
-      }
+        var mappedPort = testcontainers.GetMappedPublicPort(this.Options.Port);
+        this.Options.Port = mappedPort;
+        var client = new HttpClient();
+        var message = new HttpRequestMessage(this.Options.Method, this.Options.Uri);
+        if (this.Options.RequestContent is not null && (this.Options.Method == HttpMethod.Post || this.Options.Method == HttpMethod.Put))
+        {
+          message.Content = this.Options.RequestContent;
+        }
 
-      var sendTask = Task.Run(async () =>
-      {
-        HttpResponseMessage response = null;
-        while (response is null || !this.Options.ExpectedResponseCodes.Contains(response.StatusCode))
+        if (this.Options.UseAuth && this.Options.AuthString is not null)
+        {
+          message.Headers.Authorization = AuthenticationHeaderValue.Parse(this.Options.AuthString.ToString());
+        }
+
+        var sendTask = Task.Run(async () =>
         {
-          response = await client.SendAsync(message);
-          if (!this.Options.ExpectedResponseCodes.Contains(response.StatusCode))
+          HttpResponseMessage response = null;
+          while (response is null || !this.Options.ExpectedResponseCodes.Contains(response.StatusCode))
           {
-            Thread.Sleep(TimeSpan.FromSeconds(this.Options.RequestDelay));
+            response = await client.SendAsync(message);
+            if (++this.RetryCount > this.Options.MaxRetries)
+            {
+              throw new TimeoutException($"Http Wait Failed {this.Options.MaxRetries} Times");
+            }
+
+            if (!this.Options.ExpectedResponseCodes.Contains(response.StatusCode))
+            {
+              Thread.Sleep(TimeSpan.FromSeconds(this.Options.RequestDelay));
+            }
           }
+          return response;
+        });
+        var completed = sendTask.Wait(this.Options.TimeOut);
+        if (!completed)
+        {
+          throw new TimeoutException($"Http Wait Failed Timed Out after {this.Options.TimeOut}");
         }
 
-        return response;
-      });
-      var completed = sendTask.Wait(this.Options.TimeOut);
-      if (!completed)
+        var responseContent = await sendTask.Result.Content.ReadAsStringAsync();
+        return !this.Options.ValidateContent || Regex.Match(this.Options.ExpectedOutput, responseContent).Success;
+      }
+      catch (Exception)
       {
+        if (++this.RetryCount > this.Options.MaxRetries)
+        {
+          throw new TimeoutException($"Http Wait Failed {this.Options.MaxRetries} Times");
+        }
+
         return false;
       }
-
-      var responseContent = await sendTask.Result.Content.ReadAsStringAsync();
-      return !this.Options.ValidateContent || Regex.Match(this.Options.ExpectedOutput, responseContent).Success;
     }
   }
 }