From 86b51f55c31332903f7c2989fb70a1710f8e3293 Mon Sep 17 00:00:00 2001
From: Martin Gaievski <gaievski@amazon.com>
Date: Thu, 13 Jul 2023 19:22:30 -0700
Subject: [PATCH] Minor refactoring in uni tests

Signed-off-by: Martin Gaievski <gaievski@amazon.com>
---
 .../neuralsearch/plugin/NeuralSearchTests.java        | 11 +++--------
 .../neuralsearch/query/OpenSearchQueryTestCase.java   |  7 +++++++
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java b/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java
index 778f0e5f2..c4b1d49f7 100644
--- a/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java
+++ b/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java
@@ -6,23 +6,21 @@
 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;
 import java.util.Optional;
 
-import org.opensearch.common.SuppressForbidden;
 import org.opensearch.ingest.Processor;
 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();
@@ -34,7 +32,6 @@ public void testQuerySpecs() {
         assertTrue(querySpecs.stream().anyMatch(spec -> HybridQueryBuilder.NAME.equals(spec.getName().getPreferredName())));
     }
 
-    @SuppressForbidden(reason = "manipulates system properties for testing")
     public void testQueryPhaseSearcher() {
         NeuralSearch plugin = new NeuralSearch();
         Optional<QueryPhaseSearcher> queryPhaseSearcher = plugin.getQueryPhaseSearcher();
@@ -42,15 +39,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
index 0ce8e8fd1..c45c4adce 100644
--- 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;
@@ -28,6 +29,7 @@
 import org.opensearch.Version;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.common.CheckedConsumer;
+import org.opensearch.common.SuppressForbidden;
 import org.opensearch.common.compress.CompressedXContent;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.xcontent.XContentFactory;
@@ -222,4 +224,9 @@ public float getMaxScore(int upTo) {
             }
         };
     }
+
+    @SuppressForbidden(reason = "manipulates system properties for testing")
+    protected static void initFeatureFlags() {
+        System.setProperty(NEURAL_SEARCH_HYBRID_SEARCH_ENABLED, "true");
+    }
 }