From 6859bed18fcccbd78caa30440dfd5bc3247e9c86 Mon Sep 17 00:00:00 2001
From: Austin Vazquez <macedonv@amazon.com>
Date: Wed, 19 Jun 2024 15:48:52 -0700
Subject: [PATCH] fix: add check to server start test

Signed-off-by: Austin Vazquez <macedonv@amazon.com>
---
 tests/run.go | 48 ++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 46 insertions(+), 2 deletions(-)

diff --git a/tests/run.go b/tests/run.go
index a26b77f..cdd8723 100644
--- a/tests/run.go
+++ b/tests/run.go
@@ -387,10 +387,54 @@ func Run(o *RunOption) {
 					io.WriteString(w, response) //nolint:errcheck,gosec // Function call in server handler for testing only.
 				})
 				hostPort := fnet.GetFreePort()
-				s := http.Server{Addr: fmt.Sprintf(":%d", hostPort), Handler: mux, ReadTimeout: 30 * time.Second}
+				serverAddress := fmt.Sprintf(":%d", hostPort)
+				s := http.Server{Addr: serverAddress, Handler: mux, ReadTimeout: 30 * time.Second}
 				go s.ListenAndServe() //nolint:errcheck // Asynchronously starting server for testing only.
 
-				time.Sleep(5 * time.Second)
+				upChan := make(chan struct{})
+				timeoutChan := make(chan struct{})
+				go func() {
+					verifyServerIsUp := func() error {
+						client := http.Client{}
+						resp, err := client.Get(serverAddress)
+						if err != nil {
+							return fmt.Errorf("server has not started: %w", err)
+						}
+						defer func() {
+							err := resp.Body.Close()
+							if err != nil {
+								fmt.Fprintf(ginkgo.GinkgoWriter, "Error closing response body: %v", err)
+							}
+						}()
+						return nil
+					}
+
+					attempt := 0
+					for {
+						select {
+						case <-timeoutChan:
+							return
+						default:
+						}
+
+						err := verifyServerIsUp()
+						if err != nil {
+							close(upChan)
+							return
+						}
+						attempt++
+						fmt.Fprintf(ginkgo.GinkgoWriter, "Try server %d: %v", attempt, err)
+						time.Sleep(5 * time.Second)
+					}
+				}()
+
+				select {
+				case <-time.After(30 * time.Second):
+					close(timeoutChan)
+					ginkgo.Fail("Server failed to start after 30 seconds")
+				case <-upChan:
+				}
+
 				ginkgo.DeferCleanup(s.Shutdown, ctx)
 				command.Run(o.BaseOpt, "run", "-d", "--name", testContainerName, "--add-host", "test-host:host-gateway",
 					localImages[amazonLinux2Image], "sleep", "infinity")