From 2e1786f7b6377b7f4cbda356bd724652301f0abf Mon Sep 17 00:00:00 2001
From: Imri Paran <imri.paran@gmail.com>
Date: Mon, 22 Apr 2024 11:10:59 +0200
Subject: [PATCH] MINOR: use polling mechanism for bnackend tests (#15962)

* handle unstable tests

* handle retry by polling ES

* remove waitForAsyncOp

* use retry

* format

* poll elastic for 10 seconds

* fixed tag delete test with polling

* poll up to 3 seconds

* fixed tag delete test again

* fixed assertEntityReferenceFromSearch

* fixed assertEntityReferenceFromSearch

* better error messages
---
 .../service/resources/EntityResourceTest.java | 112 +++++++++++-------
 .../glossary/GlossaryTermResourceTest.java    |   1 +
 .../service/util/RetryableAssertionError.java |   7 ++
 .../openmetadata/service/util/TestUtils.java  |  39 ++++++
 4 files changed, 117 insertions(+), 42 deletions(-)
 create mode 100644 openmetadata-service/src/test/java/org/openmetadata/service/util/RetryableAssertionError.java

diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/resources/EntityResourceTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/resources/EntityResourceTest.java
index 07a7f0aa2d4..9dd81c7ac13 100644
--- a/openmetadata-service/src/test/java/org/openmetadata/service/resources/EntityResourceTest.java
+++ b/openmetadata-service/src/test/java/org/openmetadata/service/resources/EntityResourceTest.java
@@ -2019,7 +2019,9 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
         permissionNotAllowed(TEST_USER_NAME, List.of(MetadataOperation.DELETE)));
   }
 
-  /** Soft delete an entity and then use restore request to restore it back */
+  /**
+   * Soft delete an entity and then use restore request to restore it back
+   */
   @Test
   @Execution(ExecutionMode.CONCURRENT)
   void delete_restore_entity_200(TestInfo test) throws IOException {
@@ -2238,25 +2240,37 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
     }
     // check if the added tag if also added in the entity in search
     assertTrue(fqnList.contains(tagLabel.getTagFQN()));
-    fqnList.clear();
     // delete the tag
     tagResourceTest.deleteEntity(tag.getId(), false, true, ADMIN_AUTH_HEADERS);
-    waitForEsAsyncOp(500);
-    response =
-        getResponseFormSearch(
-            indexMapping.getIndexName(Entity.getSearchRepository().getClusterAlias()));
-    hits = response.getHits().getHits();
-    for (SearchHit hit : hits) {
-      Map<String, Object> sourceAsMap = hit.getSourceAsMap();
-      if (sourceAsMap.get("id").toString().equals(entity.getId().toString())) {
-        @SuppressWarnings("unchecked")
-        List<Map<String, String>> listTags = (List<Map<String, String>>) sourceAsMap.get("tags");
-        listTags.forEach(tempMap -> fqnList.add(tempMap.get("tagFQN")));
-        break;
-      }
-    }
-    // check if the relationships of tag are also deleted in search
-    assertFalse(fqnList.contains(tagLabel.getTagFQN()));
+
+    T finalEntity = entity;
+    TestUtils.assertEventually(
+        test.getDisplayName(),
+        () -> {
+          fqnList.clear();
+          SearchResponse afterDeleteResponse;
+          try {
+            afterDeleteResponse =
+                getResponseFormSearch(
+                    indexMapping.getIndexName(Entity.getSearchRepository().getClusterAlias()));
+          } catch (HttpResponseException e) {
+            throw new RuntimeException(e);
+          }
+
+          SearchHit[] hitsAfterDelete = afterDeleteResponse.getHits().getHits();
+          for (SearchHit hit : hitsAfterDelete) {
+            Map<String, Object> sourceAsMap = hit.getSourceAsMap();
+            if (sourceAsMap.get("id").toString().equals(finalEntity.getId().toString())) {
+              @SuppressWarnings("unchecked")
+              List<Map<String, String>> listTags =
+                  (List<Map<String, String>>) sourceAsMap.get("tags");
+              listTags.forEach(tempMap -> fqnList.add(tempMap.get("tagFQN")));
+              break;
+            }
+          }
+          // check if the relationships of tag are also deleted in search
+          assertFalse(fqnList.contains(tagLabel.getTagFQN()));
+        });
   }
 
   @Test
@@ -2587,7 +2601,9 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
     return TestUtils.put(target, restore, entityClass, status, authHeaders);
   }
 
-  /** Helper function to create an entity, submit POST API request and validate response. */
+  /**
+   * Helper function to create an entity, submit POST API request and validate response.
+   */
   public T createAndCheckEntity(K create, Map<String, String> authHeaders) throws IOException {
     // Validate an entity that is created has all the information set in create request
     String updatedBy = SecurityUtil.getPrincipalName(authHeaders);
@@ -2712,7 +2728,9 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
     }
   }
 
-  /** Helper function to generate JSON PATCH, submit PATCH API request and validate response. */
+  /**
+   * Helper function to generate JSON PATCH, submit PATCH API request and validate response.
+   */
   protected final T patchEntityAndCheck(
       T updated,
       String originalJson,
@@ -3117,7 +3135,9 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
         .withFieldsDeleted(new ArrayList<>());
   }
 
-  /** Compare fullyQualifiedName in the entityReference */
+  /**
+   * Compare fullyQualifiedName in the entityReference
+   */
   protected static void assertReference(String expected, EntityReference actual) {
     if (expected != null) {
       assertNotNull(actual);
@@ -3128,7 +3148,9 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
     }
   }
 
-  /** Compare entity Id and types in the entityReference */
+  /**
+   * Compare entity Id and types in the entityReference
+   */
   protected static void assertReference(EntityReference expected, EntityReference actual) {
     // If the actual value is inherited, it will never match the expected
     // We just ignore the validation in these cases
@@ -3146,37 +3168,43 @@ public abstract class EntityResourceTest<T extends EntityInterface, K extends Cr
   }
 
   protected void assertEntityReferenceFromSearch(T entity, EntityReference actual)
-      throws IOException, InterruptedException {
+      throws IOException {
     RestClient searchClient = getSearchClient();
     IndexMapping index = Entity.getSearchRepository().getIndexMapping(entityType);
-    Response response;
     Request request = new Request("GET", String.format("%s/_search", index.getIndexName(null)));
     String query =
         String.format(
             "{\"query\":{\"bool\":{\"filter\":[{\"term\":{\"_id\":\"%s\"}}]}}}", entity.getId());
     request.setJsonEntity(query);
     try {
-      waitForEsAsyncOp();
-      response = searchClient.performRequest(request);
+      assertEventually(
+          "assertEntityReferenceFromSearch_" + entity.getFullyQualifiedName(),
+          () -> {
+            Response response = searchClient.performRequest(request);
+            String jsonString = EntityUtils.toString(response.getEntity());
+            @SuppressWarnings("unchecked")
+            HashMap<String, Object> map =
+                (HashMap<String, Object>) JsonUtils.readOrConvertValue(jsonString, HashMap.class);
+            @SuppressWarnings("unchecked")
+            LinkedHashMap<String, Object> hits = (LinkedHashMap<String, Object>) map.get("hits");
+            @SuppressWarnings("unchecked")
+            ArrayList<LinkedHashMap<String, Object>> hitsList =
+                (ArrayList<LinkedHashMap<String, Object>>) hits.get("hits");
+            assertEquals(1, hitsList.size());
+            LinkedHashMap<String, Object> doc = hitsList.get(0);
+            @SuppressWarnings("unchecked")
+            LinkedHashMap<String, Object> source =
+                (LinkedHashMap<String, Object>) doc.get("_source");
+
+            EntityReference domainReference =
+                JsonUtils.readOrConvertValue(source.get("domain"), EntityReference.class);
+
+            assertEquals(domainReference.getId(), actual.getId());
+            assertEquals(domainReference.getType(), actual.getType());
+          });
     } finally {
       searchClient.close();
     }
-
-    String jsonString = EntityUtils.toString(response.getEntity());
-    HashMap<String, Object> map =
-        (HashMap<String, Object>) JsonUtils.readOrConvertValue(jsonString, HashMap.class);
-    LinkedHashMap<String, Object> hits = (LinkedHashMap<String, Object>) map.get("hits");
-    ArrayList<LinkedHashMap<String, Object>> hitsList =
-        (ArrayList<LinkedHashMap<String, Object>>) hits.get("hits");
-    assertEquals(1, hitsList.size());
-    LinkedHashMap<String, Object> doc = (LinkedHashMap<String, Object>) hitsList.get(0);
-    LinkedHashMap<String, Object> source = (LinkedHashMap<String, Object>) doc.get("_source");
-
-    EntityReference domainReference =
-        JsonUtils.readOrConvertValue(source.get("domain"), EntityReference.class);
-
-    assertEquals(domainReference.getId(), actual.getId());
-    assertEquals(domainReference.getType(), actual.getType());
   }
 
   protected static void checkOwnerOwns(EntityReference owner, UUID entityId, boolean expectedOwning)
diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/resources/glossary/GlossaryTermResourceTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/resources/glossary/GlossaryTermResourceTest.java
index 47bf7be141d..7feaa91e438 100644
--- a/openmetadata-service/src/test/java/org/openmetadata/service/resources/glossary/GlossaryTermResourceTest.java
+++ b/openmetadata-service/src/test/java/org/openmetadata/service/resources/glossary/GlossaryTermResourceTest.java
@@ -215,6 +215,7 @@ public class GlossaryTermResourceTest extends EntityResourceTest<GlossaryTerm, C
             .withName("t1")
             .withGlossary(glossary.getFullyQualifiedName())
             .withDescription("desc");
+
     GlossaryTerm t1 = assertDomainInheritance(create, DOMAIN.getEntityReference());
 
     // Create terms t12 under t1 without reviewers and owner
diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/util/RetryableAssertionError.java b/openmetadata-service/src/test/java/org/openmetadata/service/util/RetryableAssertionError.java
new file mode 100644
index 00000000000..6c38d7e9966
--- /dev/null
+++ b/openmetadata-service/src/test/java/org/openmetadata/service/util/RetryableAssertionError.java
@@ -0,0 +1,7 @@
+package org.openmetadata.service.util;
+
+public class RetryableAssertionError extends Exception {
+  public RetryableAssertionError(Throwable cause) {
+    super(cause);
+  }
+}
diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/util/TestUtils.java b/openmetadata-service/src/test/java/org/openmetadata/service/util/TestUtils.java
index 1d4bd19740f..5926653e193 100644
--- a/openmetadata-service/src/test/java/org/openmetadata/service/util/TestUtils.java
+++ b/openmetadata-service/src/test/java/org/openmetadata/service/util/TestUtils.java
@@ -24,11 +24,16 @@ import static org.openmetadata.service.Entity.ADMIN_USER_NAME;
 import static org.openmetadata.service.Entity.SEPARATOR;
 import static org.openmetadata.service.security.SecurityUtil.authHeaders;
 
+import io.github.resilience4j.core.functions.CheckedRunnable;
+import io.github.resilience4j.retry.Retry;
+import io.github.resilience4j.retry.RetryConfig;
+import io.github.resilience4j.retry.RetryRegistry;
 import java.lang.reflect.Field;
 import java.net.URI;
 import java.text.DateFormat;
 import java.text.ParseException;
 import java.text.SimpleDateFormat;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.Date;
@@ -85,6 +90,7 @@ import org.openmetadata.service.resources.glossary.GlossaryTermResourceTest;
 import org.openmetadata.service.resources.tags.TagResourceTest;
 import org.openmetadata.service.resources.teams.UserResourceTest;
 import org.openmetadata.service.security.SecurityUtil;
+import org.opentest4j.AssertionFailedError;
 
 @Slf4j
 public final class TestUtils {
@@ -183,6 +189,14 @@ public final class TestUtils {
                   .withUsername("admin")
                   .withPassword("admin"));
 
+  public static RetryRegistry elasticSearchRetryRegistry =
+      RetryRegistry.of(
+          RetryConfig.custom()
+              .maxAttempts(30) // about 3 seconds
+              .waitDuration(Duration.ofMillis(100))
+              .retryExceptions(RetryableAssertionError.class)
+              .build());
+
   public static void assertCustomProperties(
       List<CustomProperty> expected, List<CustomProperty> actual) {
     if (expected == actual) { // Take care of both being null
@@ -654,4 +668,29 @@ public final class TestUtils {
     // the owner of the async operation
     TimeUnit.MILLISECONDS.sleep(milliseconds);
   }
+
+  public static void assertEventually(String name, CheckedRunnable runnable) {
+    try {
+      Retry.decorateCheckedRunnable(
+              elasticSearchRetryRegistry.retry(name),
+              () -> {
+                try {
+                  runnable.run();
+                } catch (AssertionError e) {
+                  // translate AssertionErrors to Exceptions to that retry processes
+                  // them correctly. This is required because retry library only retries on
+                  // Exceptions and:
+                  // AssertionError -> Error -> Throwable
+                  // RetryableAssertionError -> Exception -> Throwable
+                  throw new RetryableAssertionError(e);
+                }
+              })
+          .run();
+    } catch (RetryableAssertionError e) {
+      throw new AssertionFailedError(
+          "Max retries exceeded polling for eventual assert", e.getCause());
+    } catch (Throwable e) {
+      throw new AssertionFailedError("Unexpected error while running retry: " + e.getMessage(), e);
+    }
+  }
 }