test updates

datahub-project · Oct 3, 2024 · f8c52d5 · f8c52d5
1 parent ac7c9ce
commit f8c52d5
Show file tree

Hide file tree

Showing 35 changed files with 117 additions and 157 deletions.
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
@@ -636,16 +636,17 @@ public ExplainResponse explain(
     }
   }
 
-  private static void testLog(ObjectMapper mapper, SearchRequest searchRequest) {
+  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: %s",
+              "SearchRequest(indices): %s",
               mapper.writerWithDefaultPrettyPrinter().writeValueAsString(indices)));
       log.warn(
           String.format(
-              "SearchRequest: %s",
+              "SearchRequest(query): %s",
               mapper.writeValueAsString(mapper.readTree(searchRequest.source().toString()))));
     } catch (JsonProcessingException e) {
       log.warn("Error writing test log");

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/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);
   }
 }
diff --git a/...test/java/com/linkedin/metadata/search/elasticsearch/LineageServiceElasticSearchTest.java b/...test/java/com/linkedin/metadata/search/elasticsearch/LineageServiceElasticSearchTest.java
@@ -1,5 +1,7 @@
 package com.linkedin.metadata.search.elasticsearch;
 
+import static org.testng.Assert.assertNotNull;
+
 import com.linkedin.metadata.config.search.SearchConfiguration;
 import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration;
 import com.linkedin.metadata.search.LineageServiceTestBase;
@@ -12,7 +14,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;
 
 @Import({
@@ -55,14 +56,8 @@ protected SearchConfiguration getSearchConfiguration() {
     return _searchConfiguration;
   }
 
-  @NotNull
-  @Override
-  protected CustomSearchConfiguration getCustomSearchConfiguration() {
-    return _customSearchConfiguration;
-  }
-
   @Test
   public void initTest() {
-    AssertJUnit.assertNotNull(_searchClient);
+    assertNotNull(_searchClient);
   }
 }
diff --git a/...t/java/com/linkedin/metadata/search/elasticsearch/SampleDataFixtureElasticSearchTest.java b/...t/java/com/linkedin/metadata/search/elasticsearch/SampleDataFixtureElasticSearchTest.java
@@ -1,8 +1,9 @@
 package com.linkedin.metadata.search.elasticsearch;
 
-import static org.testng.AssertJUnit.assertNotNull;
+import static org.testng.Assert.assertNotNull;
 
 import com.linkedin.entity.client.EntityClient;
+import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration;
 import com.linkedin.metadata.search.SearchService;
 import com.linkedin.metadata.search.fixtures.SampleDataFixtureTestBase;
 import io.datahubproject.metadata.context.OperationContext;
@@ -37,6 +38,11 @@ public class SampleDataFixtureElasticSearchTest extends SampleDataFixtureTestBas
   @Qualifier("sampleDataOperationContext")
   protected OperationContext operationContext;
 
+  @Getter
+  @Autowired
+  @Qualifier("fixtureCustomSearchConfig")
+  protected CustomSearchConfiguration customSearchConfiguration;
+
   @Test
   public void initTest() {
     assertNotNull(searchClient);

diff --git a/.../src/test/java/com/linkedin/metadata/search/elasticsearch/SearchDAOElasticSearchTest.java b/.../src/test/java/com/linkedin/metadata/search/elasticsearch/SearchDAOElasticSearchTest.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.config.search.SearchConfiguration;
 import com.linkedin.metadata.search.query.SearchDAOTestBase;

diff --git a/.../test/java/com/linkedin/metadata/search/elasticsearch/SearchServiceElasticSearchTest.java b/.../test/java/com/linkedin/metadata/search/elasticsearch/SearchServiceElasticSearchTest.java
@@ -1,5 +1,7 @@
 package com.linkedin.metadata.search.elasticsearch;
 
+import static org.testng.Assert.assertNotNull;
+
 import com.linkedin.metadata.config.search.SearchConfiguration;
 import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration;
 import com.linkedin.metadata.search.SearchServiceTestBase;
@@ -12,7 +14,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;
 
 @Import({
@@ -55,14 +56,8 @@ protected SearchConfiguration getSearchConfiguration() {
     return _searchConfiguration;
   }
 
-  @NotNull
-  @Override
-  protected CustomSearchConfiguration getCustomSearchConfiguration() {
-    return _customSearchConfiguration;
-  }
-
   @Test
   public void initTest() {
-    AssertJUnit.assertNotNull(_searchClient);
+    assertNotNull(_searchClient);
   }
 }
diff --git a/...va/com/linkedin/metadata/search/elasticsearch/SystemMetadataServiceElasticSearchTest.java b/...va/com/linkedin/metadata/search/elasticsearch/SystemMetadataServiceElasticSearchTest.java
@@ -1,5 +1,7 @@
 package com.linkedin.metadata.search.elasticsearch;
 
+import static org.testng.Assert.assertNotNull;
+
 import com.linkedin.metadata.search.elasticsearch.indexbuilder.ESIndexBuilder;
 import com.linkedin.metadata.search.elasticsearch.update.ESBulkProcessor;
 import com.linkedin.metadata.systemmetadata.SystemMetadataServiceTestBase;
@@ -8,7 +10,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})
@@ -38,6 +39,6 @@ protected ESIndexBuilder getIndexBuilder() {
 
   @Test
   public void initTest() {
-    AssertJUnit.assertNotNull(_searchClient);
+    assertNotNull(_searchClient);
   }
 }
diff --git a/...src/test/java/com/linkedin/metadata/search/elasticsearch/TestEntityElasticSearchTest.java b/...src/test/java/com/linkedin/metadata/search/elasticsearch/TestEntityElasticSearchTest.java
@@ -1,5 +1,7 @@
 package com.linkedin.metadata.search.elasticsearch;
 
+import static org.testng.Assert.assertNotNull;
+
 import com.linkedin.metadata.config.search.SearchConfiguration;
 import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration;
 import com.linkedin.metadata.search.TestEntityTestBase;
@@ -12,7 +14,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;
 
 @Import({
@@ -55,14 +56,8 @@ protected SearchConfiguration getSearchConfiguration() {
     return searchConfiguration;
   }
 
-  @NotNull
-  @Override
-  protected CustomSearchConfiguration getCustomSearchConfiguration() {
-    return customSearchConfiguration;
-  }
-
   @Test
   public void initTest() {
-    AssertJUnit.assertNotNull(searchClient);
+    assertNotNull(searchClient);
   }
 }