From 787f635464cf34d59e451faff04ad037fb6cb6d9 Mon Sep 17 00:00:00 2001
From: Tiago Queiroz <tiago.queiroz@elastic.co>
Date: Mon, 27 Nov 2023 15:18:12 +0100
Subject: [PATCH] findESDocs ensures non-zero documents

findESDocs now ensures the function returns at least one document.

Debug logs are also removed.
---
 pkg/testing/tools/estools/elasticsearch.go | 14 ----------
 testing/integration/logs_ingestion_test.go | 31 +++++++++++++++++++---
 2 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/pkg/testing/tools/estools/elasticsearch.go b/pkg/testing/tools/estools/elasticsearch.go
index 36788dd39d7..ee9c2613d82 100644
--- a/pkg/testing/tools/estools/elasticsearch.go
+++ b/pkg/testing/tools/estools/elasticsearch.go
@@ -435,7 +435,6 @@ func CheckForErrorsInLogsWithContext(ctx context.Context, client elastictranspor
 	}
 
 	return performQueryForRawQuery(ctx, queryRaw, "logs-elastic_agent*", client)
-
 }
 
 // GetLogsForDataset returns any logs associated with the datastream
@@ -535,19 +534,6 @@ func performQueryForRawQuery(ctx context.Context, queryRaw map[string]interface{
 
 	docs, err := handleDocsResponse(res)
 
-	fmt.Println(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
-	fmt.Println("Query:")
-	query, debugErr := json.MarshalIndent(queryRaw, "|", "  ")
-	if debugErr != nil {
-		fmt.Println("Error marshalling 'queryRaw':", debugErr)
-		return docs, err
-	}
-	fmt.Println("Raw Query")
-	fmt.Println(string(query))
-	fmt.Println("Documents docs.Hits.Total.Value: ", docs.Hits.Total.Value)
-	fmt.Println("Documents len(docs.Hits.Hits): ", len(docs.Hits.Hits))
-	fmt.Println(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
-
 	return docs, err
 }
 
diff --git a/testing/integration/logs_ingestion_test.go b/testing/integration/logs_ingestion_test.go
index 9757c11db05..c8f62d4e306 100644
--- a/testing/integration/logs_ingestion_test.go
+++ b/testing/integration/logs_ingestion_test.go
@@ -124,7 +124,7 @@ func testMonitoringLogsAreShipped(
 
 	// Stage 3: Make sure there are no errors in logs
 	t.Log("Making sure there are no error logs")
-	docs = findESDocs(t, func() (estools.Documents, error) {
+	docs = queryESDocs(t, func() (estools.Documents, error) {
 		return estools.CheckForErrorsInLogs(info.ESClient, info.Namespace, []string{
 			// acceptable error messages (include reason)
 			"Error dialing dial tcp 127.0.0.1:9200: connect: connection refused", // beat is running default config before its config gets updated
@@ -164,7 +164,6 @@ func testMonitoringLogsAreShipped(
 	// this field is not mapped. There is an issue for that:
 	// https://github.com/elastic/integrations/issues/6545
 	// TODO: use runtime fields while the above issue is not resolved.
-
 	docs = findESDocs(t, func() (estools.Documents, error) {
 		return estools.GetLogsForAgentID(info.ESClient, agentID)
 	})
@@ -194,13 +193,17 @@ func testMonitoringLogsAreShipped(
 	}
 }
 
-func findESDocs(t *testing.T, findFn func() (estools.Documents, error)) estools.Documents {
+// queryESDocs runs `findFn` until it returns no error
+func queryESDocs(t *testing.T, findFn func() (estools.Documents, error)) estools.Documents {
 	var docs estools.Documents
 	require.Eventually(
 		t,
 		func() bool {
 			var err error
 			docs, err = findFn()
+			if err != nil {
+				t.Logf("got an error querying ES, retrying. Error: %s", err)
+			}
 			return err == nil
 		},
 		3*time.Minute,
@@ -210,6 +213,28 @@ func findESDocs(t *testing.T, findFn func() (estools.Documents, error)) estools.
 	return docs
 }
 
+// findESDocs runs `findFn` until at least one document is returned and there is no error
+func findESDocs(t *testing.T, findFn func() (estools.Documents, error)) estools.Documents {
+	var docs estools.Documents
+	require.Eventually(
+		t,
+		func() bool {
+			var err error
+			docs, err = findFn()
+			if err != nil {
+				t.Logf("got an error querying ES, retrying. Error: %s", err)
+				return false
+			}
+
+			return docs.Hits.Total.Value != 0
+		},
+		3*time.Minute,
+		15*time.Second,
+	)
+
+	return docs
+}
+
 func testFlattenedDatastreamFleetPolicy(
 	t *testing.T,
 	ctx context.Context,