fix(tests): fix metadata-io tests (#11530)

datahub-project · Oct 4, 2024 · 9e0f68c · 9e0f68c
1 parent 134ad21
commit 9e0f68c
Show file tree

Hide file tree

Showing 57 changed files with 540 additions and 244 deletions.
diff --git a/...ata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java b/...ata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/ElasticSearchService.java
@@ -2,6 +2,7 @@
 
 import static com.linkedin.metadata.search.utils.SearchUtils.applyDefaultSearchFlags;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.metadata.browse.BrowseResult;
 import com.linkedin.metadata.browse.BrowseResultV2;
@@ -30,6 +31,7 @@
 import java.util.Optional;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.opensearch.action.explain.ExplainResponse;
@@ -51,7 +53,7 @@ public class ElasticSearchService implements EntitySearchService, ElasticSearchI
 
   private static final int MAX_RUN_IDS_INDEXED = 25; // Save the previous 25 run ids in the index.
   private final EntityIndexBuilders indexBuilders;
-  private final ESSearchDAO esSearchDAO;
+  @VisibleForTesting @Getter private final ESSearchDAO esSearchDAO;
   private final ESBrowseDAO esBrowseDAO;
   private final ESWriteDAO esWriteDAO;
 

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java
@@ -6,7 +6,9 @@
 
 import com.codahale.metrics.Timer;
 import com.datahub.util.exception.ESQueryException;
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.annotations.VisibleForTesting;
 import com.linkedin.data.template.LongMap;
 import com.linkedin.metadata.config.search.SearchConfiguration;
@@ -78,6 +80,24 @@ public class ESSearchDAO {
   @Nonnull private final SearchConfiguration searchConfiguration;
   @Nullable private final CustomSearchConfiguration customSearchConfiguration;
   @Nonnull private final QueryFilterRewriteChain queryFilterRewriteChain;
+  private final boolean testLoggingEnabled;
+
+  public ESSearchDAO(
+      RestHighLevelClient client,
+      boolean pointInTimeCreationEnabled,
+      String elasticSearchImplementation,
+      @Nonnull SearchConfiguration searchConfiguration,
+      @Nullable CustomSearchConfiguration customSearchConfiguration,
+      @Nonnull QueryFilterRewriteChain queryFilterRewriteChain) {
+    this(
+        client,
+        pointInTimeCreationEnabled,
+        elasticSearchImplementation,
+        searchConfiguration,
+        customSearchConfiguration,
+        queryFilterRewriteChain,
+        false);
+  }
 
   public long docCount(@Nonnull OperationContext opContext, @Nonnull String entityName) {
     return docCount(opContext, entityName, null);
@@ -279,6 +299,11 @@ public SearchResult search(
     searchRequest.indices(
         entityNames.stream().map(indexConvention::getEntityIndexName).toArray(String[]::new));
     searchRequestTimer.stop();
+
+    if (testLoggingEnabled) {
+      testLog(opContext.getObjectMapper(), searchRequest);
+    }
+
     // Step 2: execute the query and extract results, validated against document model as well
     return executeAndExtract(opContext, entitySpecs, searchRequest, transformedFilters, from, size);
   }
@@ -478,6 +503,11 @@ public ScrollResult scroll(
     }
 
     scrollRequestTimer.stop();
+
+    if (testLoggingEnabled) {
+      testLog(opContext.getObjectMapper(), searchRequest);
+    }
+
     return executeAndExtract(
         opContext, entitySpecs, searchRequest, transformedFilters, keepAlive, size);
   }
@@ -605,4 +635,21 @@ public ExplainResponse explain(
       throw new IllegalStateException("Failed to explain query:", e);
     }
   }
+
+  private void testLog(ObjectMapper mapper, SearchRequest searchRequest) {
+    try {
+      log.warn("SearchRequest(custom): {}", mapper.writeValueAsString(customSearchConfiguration));
+      final String[] indices = searchRequest.indices();
+      log.warn(
+          String.format(
+              "SearchRequest(indices): %s",
+              mapper.writerWithDefaultPrettyPrinter().writeValueAsString(indices)));
+      log.warn(
+          String.format(
+              "SearchRequest(query): %s",
+              mapper.writeValueAsString(mapper.readTree(searchRequest.source().toString()))));
+    } catch (JsonProcessingException e) {
+      log.warn("Error writing test log");
+    }
+  }
 }
diff --git a/...n/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java b/...n/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java
@@ -58,6 +58,7 @@
 import java.util.stream.Stream;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.collections.CollectionUtils;
 import org.opensearch.action.search.SearchRequest;
@@ -80,7 +81,7 @@ public class SearchRequestHandler {
   private static final Map<List<EntitySpec>, SearchRequestHandler> REQUEST_HANDLER_BY_ENTITY_NAME =
       new ConcurrentHashMap<>();
   private final List<EntitySpec> entitySpecs;
-  private final Set<String> defaultQueryFieldNames;
+  @Getter private final Set<String> defaultQueryFieldNames;
   @Nonnull private final HighlightBuilder highlights;
 
   private final SearchConfiguration configs;

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/DeleteEntityServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/DeleteEntityServiceTest.java
@@ -2,7 +2,9 @@
 
 import static com.linkedin.metadata.search.utils.QueryUtils.*;
 import static org.mockito.Mockito.*;
-import static org.testng.AssertJUnit.*;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
 
 import com.datahub.util.RecordUtils;
 import com.google.common.collect.ImmutableList;

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/ebean/EbeanAspectDaoTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/ebean/EbeanAspectDaoTest.java
@@ -4,6 +4,8 @@
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertTrue;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.linkedin.metadata.EbeanTestUtils;
 import com.linkedin.metadata.aspect.batch.AspectsBatch;
 import com.linkedin.metadata.config.EbeanConfiguration;
@@ -41,15 +43,15 @@ public void testGetNextVersionForUpdate() {
     // Get the captured SQL statements
     List<String> sql =
         LoggedSql.stop().stream()
-            .filter(str -> !str.contains("INFORMATION_SCHEMA.TABLES"))
+            .filter(str -> str.contains("(t0.urn,t0.aspect,t0.version)"))
             .toList();
-    assertEquals(sql.size(), 2, String.format("Found: %s", sql));
+    assertEquals(sql.size(), 1, String.format("Found: %s", sql));
     assertTrue(
         sql.get(0).contains("for update;"), String.format("Did not find `for update` in %s ", sql));
   }
 
   @Test
-  public void testGetLatestAspectsForUpdate() {
+  public void testGetLatestAspectsForUpdate() throws JsonProcessingException {
     LoggedSql.start();
 
     testDao.runInTransactionWithRetryUnlocked(
@@ -63,9 +65,10 @@ public void testGetLatestAspectsForUpdate() {
     // Get the captured SQL statements
     List<String> sql =
         LoggedSql.stop().stream()
-            .filter(str -> !str.contains("INFORMATION_SCHEMA.TABLES"))
+            .filter(str -> str.contains("(t0.urn,t0.aspect,t0.version)"))
             .toList();
-    assertEquals(sql.size(), 1, String.format("Found: %s", sql));
+    assertEquals(
+        sql.size(), 1, String.format("Found: %s", new ObjectMapper().writeValueAsString(sql)));
     assertTrue(
         sql.get(0).contains("for update;"), String.format("Did not find `for update` in %s ", sql));
   }

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/extractor/AspectExtractorTest.java b/metadata-io/src/test/java/com/linkedin/metadata/extractor/AspectExtractorTest.java
@@ -1,6 +1,6 @@
 package com.linkedin.metadata.extractor;
 
-import static org.testng.AssertJUnit.assertEquals;
+import static org.testng.Assert.assertEquals;
 
 import com.datahub.test.TestEntityAspect;
 import com.datahub.test.TestEntityAspectArray;

diff --git a/...com/linkedin/metadata/graph/search/elasticsearch/SearchGraphServiceElasticSearchTest.java b/...com/linkedin/metadata/graph/search/elasticsearch/SearchGraphServiceElasticSearchTest.java
@@ -1,5 +1,7 @@
 package com.linkedin.metadata.graph.search.elasticsearch;
 
+import static org.testng.Assert.assertNotNull;
+
 import com.linkedin.metadata.graph.search.SearchGraphServiceTestBase;
 import com.linkedin.metadata.search.elasticsearch.ElasticSearchSuite;
 import com.linkedin.metadata.search.elasticsearch.indexbuilder.ESIndexBuilder;
@@ -9,7 +11,6 @@
 import org.opensearch.client.RestHighLevelClient;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Import;
-import org.testng.AssertJUnit;
 import org.testng.annotations.Test;
 
 @Import({ElasticSearchSuite.class, SearchTestContainerConfiguration.class})
@@ -39,6 +40,6 @@ protected ESIndexBuilder getIndexBuilder() {
 
   @Test
   public void initTest() {
-    AssertJUnit.assertNotNull(_searchClient);
+    assertNotNull(_searchClient);
   }
 }
diff --git a/.../java/com/linkedin/metadata/graph/search/opensearch/SearchGraphServiceOpenSearchTest.java b/.../java/com/linkedin/metadata/graph/search/opensearch/SearchGraphServiceOpenSearchTest.java
@@ -1,5 +1,7 @@
 package com.linkedin.metadata.graph.search.opensearch;
 
+import static org.testng.Assert.assertNotNull;
+
 import com.linkedin.metadata.graph.search.SearchGraphServiceTestBase;
 import com.linkedin.metadata.search.elasticsearch.indexbuilder.ESIndexBuilder;
 import com.linkedin.metadata.search.elasticsearch.update.ESBulkProcessor;
@@ -9,7 +11,6 @@
 import org.opensearch.client.RestHighLevelClient;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Import;
-import org.testng.AssertJUnit;
 import org.testng.annotations.Test;
 
 @Import({OpenSearchSuite.class, SearchTestContainerConfiguration.class})
@@ -39,6 +40,6 @@ protected ESIndexBuilder getIndexBuilder() {
 
   @Test
   public void initTest() {
-    AssertJUnit.assertNotNull(_searchClient);
+    assertNotNull(_searchClient);
   }
 }
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/LineageSearchResultCacheKeyTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/LineageSearchResultCacheKeyTest.java
@@ -1,7 +1,7 @@
 package com.linkedin.metadata.search;
 
-import static org.testng.AssertJUnit.assertEquals;
-import static org.testng.AssertJUnit.assertNotSame;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotSame;
 
 import org.springframework.test.context.testng.AbstractTestNGSpringContextTests;
 import org.testng.annotations.Test;

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/LineageServiceTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/LineageServiceTestBase.java
@@ -35,7 +35,6 @@
 import com.linkedin.metadata.config.cache.SearchCacheConfiguration;
 import com.linkedin.metadata.config.cache.SearchLineageCacheConfiguration;
 import com.linkedin.metadata.config.search.SearchConfiguration;
-import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration;
 import com.linkedin.metadata.graph.EntityLineageResult;
 import com.linkedin.metadata.graph.GraphService;
 import com.linkedin.metadata.graph.LineageDirection;
@@ -103,9 +102,6 @@ public abstract class LineageServiceTestBase extends AbstractTestNGSpringContext
   @Nonnull
   protected abstract SearchConfiguration getSearchConfiguration();
 
-  @Nonnull
-  protected abstract CustomSearchConfiguration getCustomSearchConfiguration();
-
   private SettingsBuilder settingsBuilder;
   private ElasticSearchService elasticSearchService;
   private GraphService graphService;
@@ -211,10 +207,7 @@ private ElasticSearchService buildEntitySearchService() {
             QueryFilterRewriteChain.EMPTY);
     ESBrowseDAO browseDAO =
         new ESBrowseDAO(
-            searchClientSpy,
-            getSearchConfiguration(),
-            getCustomSearchConfiguration(),
-            QueryFilterRewriteChain.EMPTY);
+            searchClientSpy, getSearchConfiguration(), null, QueryFilterRewriteChain.EMPTY);
     ESWriteDAO writeDAO = new ESWriteDAO(searchClientSpy, getBulkProcessor(), 1);
     return new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO);
   }

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/SearchServiceTestBase.java
@@ -14,7 +14,6 @@
 import com.linkedin.common.urn.Urn;
 import com.linkedin.metadata.config.cache.EntityDocCountCacheConfiguration;
 import com.linkedin.metadata.config.search.SearchConfiguration;
-import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration;
 import com.linkedin.metadata.models.registry.SnapshotEntityRegistry;
 import com.linkedin.metadata.query.filter.Condition;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
@@ -64,9 +63,6 @@ public abstract class SearchServiceTestBase extends AbstractTestNGSpringContextT
   @Nonnull
   protected abstract SearchConfiguration getSearchConfiguration();
 
-  @Nonnull
-  protected abstract CustomSearchConfiguration getCustomSearchConfiguration();
-
   protected OperationContext operationContext;
   private SettingsBuilder settingsBuilder;
   private ElasticSearchService elasticSearchService;
@@ -136,10 +132,7 @@ private ElasticSearchService buildEntitySearchService() {
             QueryFilterRewriteChain.EMPTY);
     ESBrowseDAO browseDAO =
         new ESBrowseDAO(
-            getSearchClient(),
-            getSearchConfiguration(),
-            getCustomSearchConfiguration(),
-            QueryFilterRewriteChain.EMPTY);
+            getSearchClient(), getSearchConfiguration(), null, QueryFilterRewriteChain.EMPTY);
     ESWriteDAO writeDAO = new ESWriteDAO(getSearchClient(), getBulkProcessor(), 1);
     return new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO);
   }

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/TestEntityTestBase.java b/metadata-io/src/test/java/com/linkedin/metadata/search/TestEntityTestBase.java
@@ -13,7 +13,6 @@
 import com.linkedin.metadata.browse.BrowseResult;
 import com.linkedin.metadata.browse.BrowseResultV2;
 import com.linkedin.metadata.config.search.SearchConfiguration;
-import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration;
 import com.linkedin.metadata.models.registry.SnapshotEntityRegistry;
 import com.linkedin.metadata.search.elasticsearch.ElasticSearchService;
 import com.linkedin.metadata.search.elasticsearch.indexbuilder.ESIndexBuilder;
@@ -53,9 +52,6 @@ public abstract class TestEntityTestBase extends AbstractTestNGSpringContextTest
   @Nonnull
   protected abstract SearchConfiguration getSearchConfiguration();
 
-  @Nonnull
-  protected abstract CustomSearchConfiguration getCustomSearchConfiguration();
-
   private SettingsBuilder settingsBuilder;
   private ElasticSearchService elasticSearchService;
   private OperationContext opContext;
@@ -102,10 +98,7 @@ private ElasticSearchService buildService() {
             QueryFilterRewriteChain.EMPTY);
     ESBrowseDAO browseDAO =
         new ESBrowseDAO(
-            getSearchClient(),
-            getSearchConfiguration(),
-            getCustomSearchConfiguration(),
-            QueryFilterRewriteChain.EMPTY);
+            getSearchClient(), getSearchConfiguration(), null, QueryFilterRewriteChain.EMPTY);
     ESWriteDAO writeDAO = new ESWriteDAO(getSearchClient(), getBulkProcessor(), 1);
     ElasticSearchService searchService =
         new ElasticSearchService(indexBuilders, searchDAO, browseDAO, writeDAO);

diff --git a/...-io/src/test/java/com/linkedin/metadata/search/elasticsearch/GoldenElasticSearchTest.java b/...-io/src/test/java/com/linkedin/metadata/search/elasticsearch/GoldenElasticSearchTest.java
@@ -1,6 +1,6 @@
 package com.linkedin.metadata.search.elasticsearch;
 
-import static org.testng.AssertJUnit.assertNotNull;
+import static org.testng.Assert.assertNotNull;
 
 import com.linkedin.metadata.search.SearchService;
 import com.linkedin.metadata.search.fixtures.GoldenTestBase;

diff --git a/...c/test/java/com/linkedin/metadata/search/elasticsearch/IndexBuilderElasticSearchTest.java b/...c/test/java/com/linkedin/metadata/search/elasticsearch/IndexBuilderElasticSearchTest.java
@@ -1,6 +1,6 @@
 package com.linkedin.metadata.search.elasticsearch;
 
-import static org.testng.AssertJUnit.assertNotNull;
+import static org.testng.Assert.assertNotNull;
 
 import com.linkedin.metadata.search.indexbuilder.IndexBuilderTestBase;
 import io.datahubproject.test.search.config.SearchTestContainerConfiguration;

diff --git a/.../java/com/linkedin/metadata/search/elasticsearch/LineageDataFixtureElasticSearchTest.java b/.../java/com/linkedin/metadata/search/elasticsearch/LineageDataFixtureElasticSearchTest.java
@@ -1,5 +1,7 @@
 package com.linkedin.metadata.search.elasticsearch;
 
+import static org.testng.Assert.assertNotNull;
+
 import com.linkedin.metadata.search.LineageSearchService;
 import com.linkedin.metadata.search.SearchService;
 import com.linkedin.metadata.search.fixtures.LineageDataFixtureTestBase;
@@ -10,7 +12,6 @@
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.context.annotation.Import;
-import org.testng.AssertJUnit;
 import org.testng.annotations.Test;
 
 @Getter
@@ -35,6 +36,6 @@ public class LineageDataFixtureElasticSearchTest extends LineageDataFixtureTestB
 
   @Test
   public void initTest() {
-    AssertJUnit.assertNotNull(lineageService);
+    assertNotNull(lineageService);
   }
 }