Minor refactoring in uni tests

Signed-off-by: Martin Gaievski <gaievski@amazon.com>
opensearch-project · Jul 14, 2023 · 267a716 · 267a716
1 parent a6c52bb
commit 267a716
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java b/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java
@@ -6,7 +6,6 @@
 package org.opensearch.neuralsearch.plugin;
 
 import static org.mockito.Mockito.mock;
-import static org.opensearch.neuralsearch.plugin.NeuralSearch.NEURAL_SEARCH_HYBRID_SEARCH_ENABLED;
 
 import java.util.List;
 import java.util.Map;
@@ -17,12 +16,12 @@
 import org.opensearch.neuralsearch.processor.TextEmbeddingProcessor;
 import org.opensearch.neuralsearch.query.HybridQueryBuilder;
 import org.opensearch.neuralsearch.query.NeuralQueryBuilder;
+import org.opensearch.neuralsearch.query.OpenSearchQueryTestCase;
 import org.opensearch.neuralsearch.search.query.HybridQueryPhaseSearcher;
 import org.opensearch.plugins.SearchPlugin;
 import org.opensearch.search.query.QueryPhaseSearcher;
-import org.opensearch.test.OpenSearchTestCase;
 
-public class NeuralSearchTests extends OpenSearchTestCase {
+public class NeuralSearchTests extends OpenSearchQueryTestCase {
 
     public void testQuerySpecs() {
         NeuralSearch plugin = new NeuralSearch();
@@ -42,15 +41,13 @@ public void testQueryPhaseSearcher() {
         assertNotNull(queryPhaseSearcher);
         assertTrue(queryPhaseSearcher.isEmpty());
 
-        System.setProperty(NEURAL_SEARCH_HYBRID_SEARCH_ENABLED, "true");
+        initFeatureFlags();
 
         Optional<QueryPhaseSearcher> queryPhaseSearcherWithFeatureFlagDisabled = plugin.getQueryPhaseSearcher();
 
         assertNotNull(queryPhaseSearcherWithFeatureFlagDisabled);
         assertFalse(queryPhaseSearcherWithFeatureFlagDisabled.isEmpty());
         assertTrue(queryPhaseSearcherWithFeatureFlagDisabled.get() instanceof HybridQueryPhaseSearcher);
-
-        System.setProperty(NEURAL_SEARCH_HYBRID_SEARCH_ENABLED, "");
     }
 
     public void testProcessors() {

diff --git a/src/test/java/org/opensearch/neuralsearch/query/OpenSearchQueryTestCase.java b/src/test/java/org/opensearch/neuralsearch/query/OpenSearchQueryTestCase.java
@@ -8,6 +8,7 @@
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.singletonMap;
 import static java.util.stream.Collectors.toList;
+import static org.opensearch.neuralsearch.plugin.NeuralSearch.NEURAL_SEARCH_HYBRID_SEARCH_ENABLED;
 
 import java.io.IOException;
 import java.util.Arrays;
@@ -222,4 +223,8 @@ public float getMaxScore(int upTo) {
             }
         };
     }
+
+    protected static void initFeatureFlags() {
+        System.setProperty(NEURAL_SEARCH_HYBRID_SEARCH_ENABLED, "true");
+    }
 }