fix tests

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")

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.