From a8f6bc2d5aae5564c38b8fd368aa84e08e201c0c Mon Sep 17 00:00:00 2001
From: Kevin Park <krapi0314@gmail.com>
Date: Thu, 8 Feb 2024 13:55:10 +0900
Subject: [PATCH] Add Test Server Wait Helper Function (#787)

Co-authored-by: Youngteac Hong <susukang98@gmail.com>
---
 server/rpc/server_test.go |  5 ++++-
 test/helper/helper.go     | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/server/rpc/server_test.go b/server/rpc/server_test.go
index 745717ad8..8c3d3ff3a 100644
--- a/server/rpc/server_test.go
+++ b/server/rpc/server_test.go
@@ -105,9 +105,12 @@ func TestMain(m *testing.M) {
 		log.Fatal(err)
 	}
 
-	if err := testRPCServer.Start(); err != nil {
+	if err = testRPCServer.Start(); err != nil {
 		log.Fatalf("failed rpc listen: %s\n", err)
 	}
+	if err = helper.WaitForServerToStart(testRPCAddr); err != nil {
+		log.Fatal(err)
+	}
 
 	authInterceptor := client.NewAuthInterceptor(project.PublicKey, "")
 
diff --git a/test/helper/helper.go b/test/helper/helper.go
index 362270f06..1dc282b74 100644
--- a/test/helper/helper.go
+++ b/test/helper/helper.go
@@ -21,6 +21,7 @@ import (
 	"context"
 	"fmt"
 	"log"
+	"net"
 	"strings"
 	"testing"
 	gotime "time"
@@ -454,3 +455,34 @@ func CreateDummyClientWithID(
 
 	return nil
 }
+
+// WaitForServerToStart waits for the server to start.
+func WaitForServerToStart(addr string) error {
+	maxRetries := 10
+	initialDelay := 100 * gotime.Millisecond
+	maxDelay := 5 * gotime.Second
+
+	for attempt := 0; attempt < maxRetries; attempt++ {
+		// Exponential backoff calculation
+		delay := initialDelay * gotime.Duration(1<<uint(attempt))
+		fmt.Println("delay: ", delay)
+		if delay > maxDelay {
+			delay = maxDelay
+		}
+
+		conn, err := net.DialTimeout("tcp", addr, 1*gotime.Second)
+		if err != nil {
+			gotime.Sleep(delay)
+			continue
+		}
+
+		err = conn.Close()
+		if err != nil {
+			return fmt.Errorf("close connection: %w", err)
+		}
+
+		return nil
+	}
+
+	return fmt.Errorf("timeout for server to start: %s", addr)
+}