From 362fa95bb265861dc0bd7420e58c17a0d3c3dc1a Mon Sep 17 00:00:00 2001
From: vzaytsev <victor.zaytsev@h2o.ai>
Date: Tue, 12 Dec 2023 11:26:36 -0500
Subject: [PATCH] fix tests

---
 metadata.go   | 4 ++--
 tasks_test.go | 4 +++-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/metadata.go b/metadata.go
index e3f9fd4..6b2cd81 100644
--- a/metadata.go
+++ b/metadata.go
@@ -79,7 +79,7 @@ func (md *MetadataService) AwaitDataSourceAvailable(dataSourceName string) error
 				return err
 			}
 			if len(res) >= 1 && res[0].Cnt == 1 {
-				break
+				return nil
 			}
 		case <-time.After(md.awaitTimeout):
 			return errors.New("AwaitDataSourceAvailable timeout")
@@ -106,7 +106,7 @@ func (md *MetadataService) AwaitRecordsCount(dataSourceName string, recordsCount
 			}
 
 			if len(res) >= 1 && res[0].Cnt == recordsCount {
-				break
+				return nil
 			}
 		case <-time.After(md.awaitTimeout):
 			return errors.New("AwaitRecordsCount timeout")
diff --git a/tasks_test.go b/tasks_test.go
index 1df29b2..653b11a 100644
--- a/tasks_test.go
+++ b/tasks_test.go
@@ -59,6 +59,7 @@ func triggerIngestionTask(d *Client, dataSourceName string, entries []testDO) (s
 func AwaitTaskCompletion(client *Client, taskID string, awaitTimeout time.Duration, tickerDuration time.Duration) error {
 	ticker := time.NewTicker(tickerDuration)
 	defer ticker.Stop()
+L:
 	for {
 		select {
 		case <-ticker.C:
@@ -70,11 +71,12 @@ func AwaitTaskCompletion(client *Client, taskID string, awaitTimeout time.Durati
 			if res.Status.Status == "RUNNING" {
 				continue
 			}
-			break
+			break L
 		case <-time.After(awaitTimeout):
 			return errors.New("AwaitTaskRunning timeout")
 		}
 	}
+	return nil
 }
 
 // AwaitTaskStatus waits for the druid task status for the maximum of awaitTimeout duration, querying druid task API.