From 55d48b79b5b859e6622c318fd8c9061cbcb07732 Mon Sep 17 00:00:00 2001
From: Yury Frolov <yura200253@gmail.com>
Date: Mon, 15 Apr 2024 14:36:27 +0500
Subject: [PATCH] Rewrite health check

---
 test/feature/features/scram_auth.feature | 20 ++++----------------
 test/feature/spqr_test.go                | 15 ++++++++++++++-
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/test/feature/features/scram_auth.feature b/test/feature/features/scram_auth.feature
index de196f0d2..ee5ed1af7 100644
--- a/test/feature/features/scram_auth.feature
+++ b/test/feature/features/scram_auth.feature
@@ -22,32 +22,20 @@ Feature: SCRAM auth test
     ROUTER_CONFIG=/spqr/test/feature/conf/router_with_scram_backend.yaml
     """
     Given cluster is up and running
-    When I run command on host "shard1" with timeout "220" seconds
+    When I run command on host "shard1"
     """
     echo 'host all all all scram-sha-256' > /var/lib/postgresql/13/main/pg_hba.conf
     service postgresql reload
-    for (( i = 0; i < 10; i++)); do
-        if PGPASSWORD=12345678 psql -h spqr_shard_1 -p 6432 -d db1 -U user1 -c "SELECT 1"; then
-            break
-        else
-            sleep 20
-        fi
-    done
     """
     Then command return code should be "0"
-    When I run command on host "shard2" with timeout "220" seconds
+    And I wait for host "shard1" to respond
+    When I run command on host "shard2"
     """
     echo 'host all all all scram-sha-256' > /var/lib/postgresql/13/main/pg_hba.conf
     service postgresql reload
-    for (( i = 0; i < 10; i++)); do
-        if PGPASSWORD=12345678 psql -h spqr_shard_2 -p 6432 -d db1 -U user1 -c "SELECT 1"; then
-            break
-        else
-            sleep 20
-        fi
-    done
     """
     Then command return code should be "0"
+    And I wait for host "shard2" to respond
     When I run SQL on host "router"
     """
     SELECT 1
diff --git a/test/feature/spqr_test.go b/test/feature/spqr_test.go
index 025cc21d4..c4bd08c2a 100644
--- a/test/feature/spqr_test.go
+++ b/test/feature/spqr_test.go
@@ -649,6 +649,19 @@ func (tctx *testContext) stepHostIsStarted(service string) error {
 	return fmt.Errorf("service %s was not found in docker composer", service)
 }
 
+func (tctx *testContext) stepWaitPostgresqlToRespond(host string) error {
+	const trials = 10
+	const timeout = 20 * time.Second
+	for i := 0; i < trials; i++ {
+		_, err := tctx.queryPostgresql(host, "SELECT 1", struct{}{})
+		if err == nil {
+			return nil
+		}
+		time.Sleep(timeout)
+	}
+	return fmt.Errorf("host \"%s\" did not respond until timeout", host)
+}
+
 func (tctx *testContext) stepIRunCommandOnHost(host string, body *godog.DocString) error {
 	cmd := strings.TrimSpace(body.Content)
 	var err error
@@ -940,11 +953,11 @@ func InitializeScenario(s *godog.ScenarioContext, t *testing.T, debug bool) {
 	s.Step(`^SQL error on host "([^"]*)" should match (\w+)$`, tctx.stepErrorShouldMatch)
 	s.Step(`^file "([^"]*)" on host "([^"]*)" should match (\w+)$`, tctx.stepFileOnHostShouldMatch)
 	s.Step(`^I fail to run SQL on host "([^"]*)"$`, tctx.stepIFailSQLOnHost)
+	s.Step(`^I wait for host "([^"]*)" to respond$`, tctx.stepWaitPostgresqlToRespond)
 
 	// variable manipulation
 	s.Step(`^we save response row "([^"]*)" column "([^"]*)"$`, tctx.stepSaveResponseBodyAtPathAsJSON)
 	s.Step(`^hide "([^"]*)" field$`, tctx.stepHideField)
-
 }
 
 func TestSpqr(t *testing.T) {