move common feature disablement code to util class

Signed-off-by: HenryL27 <[email protected]>
HenryL27 · Feb 13, 2024 · e1bd813 · e1bd813
1 parent b3936d9
commit e1bd813
Show file tree

Hide file tree

Showing 10 changed files with 59 additions and 36 deletions.
diff --git a/memory/src/test/java/org/opensearch/ml/memory/MemoryTestUtil.java b/memory/src/test/java/org/opensearch/ml/memory/MemoryTestUtil.java
@@ -0,0 +1,41 @@
+/*
+ * Copyright 2024 Aryn
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.opensearch.ml.memory;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.util.Set;
+
+import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.settings.ClusterSettings;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.ml.common.conversation.ConversationalIndexConstants;
+
+public class MemoryTestUtil {
+
+    public static ClusterService clusterServiceWithMemoryFeatureDisabled() {
+        ClusterService mockClusterService = mock(ClusterService.class);
+        Settings settings = Settings.builder().put(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED.getKey(), false).build();
+        when(mockClusterService.getSettings()).thenReturn(settings);
+        when(mockClusterService.getClusterSettings())
+            .thenReturn(new ClusterSettings(settings, Set.of(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED)));
+        return mockClusterService;
+    }
+
+}
diff --git a/.../org/opensearch/ml/memory/action/conversation/CreateConversationTransportActionTests.java b/.../org/opensearch/ml/memory/action/conversation/CreateConversationTransportActionTests.java
@@ -41,6 +41,7 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.ml.common.conversation.ConversationalIndexConstants;
+import org.opensearch.ml.memory.MemoryTestUtil;
 import org.opensearch.ml.memory.index.OpenSearchConversationalMemoryHandler;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
@@ -154,10 +155,7 @@ public void testDoExecuteFails_thenFail() {
     }
 
     public void testFeatureDisabled_ThenFail() {
-        Settings settings = Settings.builder().put(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED.getKey(), false).build();
-        when(this.clusterService.getSettings()).thenReturn(settings);
-        when(this.clusterService.getClusterSettings())
-            .thenReturn(new ClusterSettings(settings, Set.of(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED)));
+        clusterService = MemoryTestUtil.clusterServiceWithMemoryFeatureDisabled();
         this.action = spy(new CreateConversationTransportAction(transportService, actionFilters, cmHandler, client, clusterService));
 
         action.doExecute(null, request, actionListener);

diff --git a/...a/org/opensearch/ml/memory/action/conversation/CreateInteractionTransportActionTests.java b/...a/org/opensearch/ml/memory/action/conversation/CreateInteractionTransportActionTests.java
@@ -45,6 +45,7 @@
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.ml.common.conversation.ConversationalIndexConstants;
+import org.opensearch.ml.memory.MemoryTestUtil;
 import org.opensearch.ml.memory.index.OpenSearchConversationalMemoryHandler;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
@@ -227,10 +228,7 @@ public void testDoExecuteFails_thenFail() {
     }
 
     public void testFeatureDisabled_ThenFail() {
-        Settings settings = Settings.builder().put(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED.getKey(), false).build();
-        when(this.clusterService.getSettings()).thenReturn(settings);
-        when(this.clusterService.getClusterSettings())
-            .thenReturn(new ClusterSettings(settings, Set.of(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED)));
+        clusterService = MemoryTestUtil.clusterServiceWithMemoryFeatureDisabled();
         this.action = spy(new CreateInteractionTransportAction(transportService, actionFilters, cmHandler, client, clusterService));
 
         action.doExecute(null, request, actionListener);

diff --git a/.../org/opensearch/ml/memory/action/conversation/DeleteConversationTransportActionTests.java b/.../org/opensearch/ml/memory/action/conversation/DeleteConversationTransportActionTests.java
@@ -40,6 +40,7 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.ml.common.conversation.ConversationalIndexConstants;
+import org.opensearch.ml.memory.MemoryTestUtil;
 import org.opensearch.ml.memory.index.OpenSearchConversationalMemoryHandler;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
@@ -138,10 +139,7 @@ public void testdoExecuteFails_thenFail() {
     }
 
     public void testFeatureDisabled_ThenFail() {
-        Settings settings = Settings.builder().put(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED.getKey(), false).build();
-        when(this.clusterService.getSettings()).thenReturn(settings);
-        when(this.clusterService.getClusterSettings())
-            .thenReturn(new ClusterSettings(settings, Set.of(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED)));
+        clusterService = MemoryTestUtil.clusterServiceWithMemoryFeatureDisabled();
         this.action = spy(new DeleteConversationTransportAction(transportService, actionFilters, cmHandler, client, clusterService));
 
         action.doExecute(null, request, actionListener);

diff --git a/...ava/org/opensearch/ml/memory/action/conversation/GetConversationTransportActionTests.java b/...ava/org/opensearch/ml/memory/action/conversation/GetConversationTransportActionTests.java
@@ -43,6 +43,7 @@
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.ml.common.conversation.ConversationMeta;
 import org.opensearch.ml.common.conversation.ConversationalIndexConstants;
+import org.opensearch.ml.memory.MemoryTestUtil;
 import org.opensearch.ml.memory.index.OpenSearchConversationalMemoryHandler;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
@@ -138,10 +139,7 @@ public void testHandlerThrows_ThenFail() {
     }
 
     public void testFeatureDisabled_ThenFail() {
-        Settings settings = Settings.builder().put(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED.getKey(), false).build();
-        when(this.clusterService.getSettings()).thenReturn(settings);
-        when(this.clusterService.getClusterSettings())
-            .thenReturn(new ClusterSettings(settings, Set.of(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED)));
+        clusterService = MemoryTestUtil.clusterServiceWithMemoryFeatureDisabled();
         this.action = spy(new GetConversationTransportAction(transportService, actionFilters, cmHandler, client, clusterService));
 
         action.doExecute(null, request, actionListener);

diff --git a/...va/org/opensearch/ml/memory/action/conversation/GetConversationsTransportActionTests.java b/...va/org/opensearch/ml/memory/action/conversation/GetConversationsTransportActionTests.java
@@ -44,6 +44,7 @@
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.ml.common.conversation.ConversationMeta;
 import org.opensearch.ml.common.conversation.ConversationalIndexConstants;
+import org.opensearch.ml.memory.MemoryTestUtil;
 import org.opensearch.ml.memory.index.OpenSearchConversationalMemoryHandler;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
@@ -193,10 +194,7 @@ public void testdoExecuteFails_thenFail() {
     }
 
     public void testFeatureDisabled_ThenFail() {
-        Settings settings = Settings.builder().put(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED.getKey(), false).build();
-        when(this.clusterService.getSettings()).thenReturn(settings);
-        when(this.clusterService.getClusterSettings())
-            .thenReturn(new ClusterSettings(settings, Set.of(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED)));
+        clusterService = MemoryTestUtil.clusterServiceWithMemoryFeatureDisabled();
         this.action = spy(new GetConversationsTransportAction(transportService, actionFilters, cmHandler, client, clusterService));
 
         action.doExecute(null, request, actionListener);

diff --git a/...java/org/opensearch/ml/memory/action/conversation/GetInteractionTransportActionTests.java b/...java/org/opensearch/ml/memory/action/conversation/GetInteractionTransportActionTests.java
@@ -44,6 +44,7 @@
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.ml.common.conversation.ConversationalIndexConstants;
 import org.opensearch.ml.common.conversation.Interaction;
+import org.opensearch.ml.memory.MemoryTestUtil;
 import org.opensearch.ml.memory.index.OpenSearchConversationalMemoryHandler;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
@@ -147,10 +148,7 @@ public void testHandlerThrows_ThenFail() {
     }
 
     public void testFeatureDisabled_ThenFail() {
-        Settings settings = Settings.builder().put(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED.getKey(), false).build();
-        when(this.clusterService.getSettings()).thenReturn(settings);
-        when(this.clusterService.getClusterSettings())
-            .thenReturn(new ClusterSettings(settings, Set.of(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED)));
+        clusterService = MemoryTestUtil.clusterServiceWithMemoryFeatureDisabled();
         this.action = spy(new GetInteractionTransportAction(transportService, actionFilters, cmHandler, client, clusterService));
 
         action.doExecute(null, request, actionListener);

diff --git a/...ava/org/opensearch/ml/memory/action/conversation/GetInteractionsTransportActionTests.java b/...ava/org/opensearch/ml/memory/action/conversation/GetInteractionsTransportActionTests.java
@@ -45,6 +45,7 @@
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.ml.common.conversation.ConversationalIndexConstants;
 import org.opensearch.ml.common.conversation.Interaction;
+import org.opensearch.ml.memory.MemoryTestUtil;
 import org.opensearch.ml.memory.index.OpenSearchConversationalMemoryHandler;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
@@ -185,10 +186,7 @@ public void testDoExecuteFails_thenFail() {
     }
 
     public void testFeatureDisabled_ThenFail() {
-        Settings settings = Settings.builder().put(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED.getKey(), false).build();
-        when(this.clusterService.getSettings()).thenReturn(settings);
-        when(this.clusterService.getClusterSettings())
-            .thenReturn(new ClusterSettings(settings, Set.of(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED)));
+        clusterService = MemoryTestUtil.clusterServiceWithMemoryFeatureDisabled();
         this.action = spy(new GetInteractionsTransportAction(transportService, actionFilters, cmHandler, client, clusterService));
 
         action.doExecute(null, request, actionListener);

diff --git a/...org/opensearch/ml/memory/action/conversation/SearchConversationsTransportActionTests.java b/...org/opensearch/ml/memory/action/conversation/SearchConversationsTransportActionTests.java
@@ -43,6 +43,7 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.ml.common.conversation.ConversationalIndexConstants;
+import org.opensearch.ml.memory.MemoryTestUtil;
 import org.opensearch.ml.memory.index.OpenSearchConversationalMemoryHandler;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
@@ -109,10 +110,7 @@ public void testEnabled_ThenSucceed() {
     }
 
     public void testDisabled_ThenFail() {
-        Settings settings = Settings.builder().put(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED.getKey(), false).build();
-        when(this.clusterService.getSettings()).thenReturn(settings);
-        when(this.clusterService.getClusterSettings())
-            .thenReturn(new ClusterSettings(settings, Set.of(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED)));
+        clusterService = MemoryTestUtil.clusterServiceWithMemoryFeatureDisabled();
         this.action = spy(new SearchConversationsTransportAction(transportService, actionFilters, cmHandler, client, clusterService));
 
         action.doExecute(null, request, actionListener);

diff --git a/...org/opensearch/ml/memory/action/conversation/SearchInteractionsTransportActionsTests.java b/...org/opensearch/ml/memory/action/conversation/SearchInteractionsTransportActionsTests.java
@@ -42,6 +42,7 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.ml.common.conversation.ConversationalIndexConstants;
+import org.opensearch.ml.memory.MemoryTestUtil;
 import org.opensearch.ml.memory.index.OpenSearchConversationalMemoryHandler;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
@@ -109,10 +110,7 @@ public void testFeatureEnabled_ThenSucceed() {
     }
 
     public void testDisabled_ThenFail() {
-        Settings settings = Settings.builder().put(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED.getKey(), false).build();
-        when(this.clusterService.getSettings()).thenReturn(settings);
-        when(this.clusterService.getClusterSettings())
-            .thenReturn(new ClusterSettings(settings, Set.of(ConversationalIndexConstants.ML_COMMONS_MEMORY_FEATURE_ENABLED)));
+        clusterService = MemoryTestUtil.clusterServiceWithMemoryFeatureDisabled();
         this.action = spy(new SearchInteractionsTransportAction(transportService, actionFilters, cmHandler, client, clusterService));
 
         action.doExecute(null, request, actionListener);