add multitenant label tests

Signed-off-by: Kaushal Kumar <[email protected]>

server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java

@@ -841,7 +841,6 @@ public final ShardSearchRequest buildShardSearchRequest(SearchShardIterator shar
         // than creating an empty response in the search thread pool.
         // Note that, we have to disable this shortcut for queries that create a context (scroll and search context).
         shardRequest.canReturnNullResponseIfMatchNoDocs(hasShardResponse.get() && shardRequest.scroll() == null);
-
         return shardRequest;
     }
 

server/src/main/java/org/opensearch/action/search/SearchRequest.java

@@ -710,6 +710,7 @@ public SearchRequest pipeline(String pipeline) {
     public String pipeline() {
         return pipeline;
     }
+
     @Override
     public SearchTask createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
         return new SearchTask(id, type, action, this::buildDescription, parentTaskId, headers, cancelAfterTimeInterval);

server/src/main/java/org/opensearch/action/search/SearchShardTask.java

@@ -37,8 +37,8 @@
 import org.opensearch.core.tasks.TaskId;
 import org.opensearch.search.fetch.ShardFetchSearchRequest;
 import org.opensearch.search.internal.ShardSearchRequest;
-import org.opensearch.search.resource_limit_group.ResourceLimitGroupTask;
 import org.opensearch.tasks.CancellableTask;
+import org.opensearch.tasks.ResourceLimitGroupTask;
 import org.opensearch.tasks.SearchBackpressureTask;
 
 import java.util.Map;

server/src/main/java/org/opensearch/action/search/SearchTask.java

@@ -35,8 +35,8 @@
 import org.opensearch.common.annotation.PublicApi;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.tasks.TaskId;
-import org.opensearch.search.resource_limit_group.ResourceLimitGroupTask;
 import org.opensearch.tasks.CancellableTask;
+import org.opensearch.tasks.ResourceLimitGroupTask;
 import org.opensearch.tasks.SearchBackpressureTask;
 
 import java.util.Map;

server/src/main/java/org/opensearch/action/search/TransportSearchAction.java

@@ -123,7 +123,8 @@
 import java.util.stream.StreamSupport;
 
 import static org.opensearch.action.admin.cluster.node.tasks.get.GetTaskAction.TASKS_ORIGIN;
-import static org.opensearch.action.search.SearchType.*;
+import static org.opensearch.action.search.SearchType.DFS_QUERY_THEN_FETCH;
+import static org.opensearch.action.search.SearchType.QUERY_THEN_FETCH;
 import static org.opensearch.search.sort.FieldSortBuilder.hasPrimaryFieldSort;
 
 /**
@@ -1110,7 +1111,7 @@ private void executeSearch(
         Map<String, Object> multiTenantLabels = searchRequest.source().multiTenantLabels();
         String tenant = NOT_PROVIDED;
         if (multiTenantLabels != null) {
-            tenant = (String) multiTenantLabels.get(MultiTenantLabel.TENANT_LABEL.name());
+            tenant = (String) multiTenantLabels.get(MultiTenantLabel.TENANT.name());
         }
         task.setResourceLimitGroupName(tenant);
 

server/src/main/java/org/opensearch/rest/action/search/RestSearchAction.java

@@ -223,7 +223,6 @@ public static void parseSearchRequest(
             );
         }
 
-
         searchRequest.setCancelAfterTimeInterval(request.paramAsTime("cancel_after_time_interval", null));
     }
 

server/src/main/java/org/opensearch/search/MultiTenantLabel.java

@@ -13,25 +13,29 @@
 import org.opensearch.core.common.io.stream.Writeable;
 
 import java.io.IOException;
-import java.util.Locale;
 
 /**
  * Enum to hold all multitenant labels in workloads
  */
 public enum MultiTenantLabel implements Writeable {
     // This label is basically used to define tenancy for multiple features e,g; Query Sandboxing, Query Insights
-    TENANT_LABEL("tenant_label");
+    TENANT("tenant");
 
     private final String value;
+
     MultiTenantLabel(String name) {
         this.value = name;
     }
 
+    public String getValue() {
+        return value;
+    }
+
     public static MultiTenantLabel fromName(String name) {
-        switch (name.toLowerCase(Locale.ROOT)) {
-            // Other cases can be added for other keys in the ENUM
-            case "tenant_label":
-                return TENANT_LABEL;
+        for (MultiTenantLabel label : values()) {
+            if (label.getValue().equalsIgnoreCase(name)) {
+                return label;
+            }
         }
         throw new IllegalArgumentException("Illegal name + " + name);
     }
@@ -40,7 +44,6 @@ public static MultiTenantLabel fromName(StreamInput in) throws IOException {
         return fromName(in.readString());
     }
 
-
     /**
      * Write this into the {@linkplain StreamOutput}.
      *

server/src/main/java/org/opensearch/search/SearchService.java

@@ -603,7 +603,7 @@ public void onFailure(Exception exc) {
     private void setTenantInTask(SearchShardTask task, ShardSearchRequest request) {
         String tenant = NOT_PROVIDED;
         if (request.source().multiTenantLabels() != null) {
-            tenant = (String) request.source().multiTenantLabels().get(MultiTenantLabel.TENANT_LABEL.name());
+            tenant = (String) request.source().multiTenantLabels().get(MultiTenantLabel.TENANT.name());
         }
         task.setResourceLimitGroupName(tenant);
     }

server/src/main/java/org/opensearch/search/builder/SearchSourceBuilder.java

@@ -77,7 +77,12 @@
 import org.opensearch.search.suggest.SuggestBuilder;
 
 import java.io.IOException;
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
 
 import static org.opensearch.index.query.AbstractQueryBuilder.parseInnerQueryBuilder;
 import static org.opensearch.search.internal.SearchContext.TRACK_TOTAL_HITS_ACCURATE;

server/src/main/java/org/opensearch/tasks/ResourceLimitGroupTask.java

@@ -0,0 +1,18 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.tasks;
+
+/**
+ * Tasks which should be grouped
+ */
+public interface ResourceLimitGroupTask {
+    void setResourceLimitGroupName(String name);
+
+    String getResourceLimitGroupName();
+}

server/src/test/java/org/opensearch/search/MultiTenantLabelTest.java

@@ -0,0 +1,44 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search;
+
+import org.opensearch.core.common.io.stream.StreamInput;
+
+import java.io.IOException;
+
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+
+public class MultiTenantLabelTest extends AbstractSearchTestCase {
+
+    public void testValidMultiTenantLabel() {
+        MultiTenantLabel label = MultiTenantLabel.fromName("tenant");
+        assertEquals(label.getValue(), "tenant");
+    }
+
+    public void testInvalidMultiTenantLabel() {
+        assertThrows(IllegalArgumentException.class, () -> MultiTenantLabel.fromName("foo"));
+    }
+
+    public void testValidMultiTenantLabelWithStreamInput() throws IOException {
+        StreamInput streamInput = mock(StreamInput.class);
+        doReturn("tenant").when(streamInput).readString();
+
+        MultiTenantLabel label = MultiTenantLabel.fromName(streamInput);
+        assertEquals(label.getValue(), "tenant");
+    }
+
+    public void testInvalidMultiTenantLabelWithStreamInput() throws IOException {
+        StreamInput streamInput = mock(StreamInput.class);
+        doReturn("foo").when(streamInput).readString();
+
+        assertThrows(IllegalArgumentException.class, () -> MultiTenantLabel.fromName(streamInput));
+    }
+
+}