From 267a716556fd55065db268a0dd684f6aa00bceb5 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           | 9 +++------
 .../neuralsearch/query/OpenSearchQueryTestCase.java      | 5 +++++
 2 files changed, 8 insertions(+), 6 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..6a2c37524 100644
--- 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
index 0ce8e8fd1..21b54a603 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;
@@ -222,4 +223,8 @@ public float getMaxScore(int upTo) {
             }
         };
     }
+
+    protected static void initFeatureFlags() {
+        System.setProperty(NEURAL_SEARCH_HYBRID_SEARCH_ENABLED, "true");
+    }
 }