add tests for custom namedWritable QueryGroupMetadata

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>
opensearch-project · Jun 25, 2024 · 4508c9d · 4508c9d
1 parent 0da3b68
commit 4508c9d
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 14 deletions.
diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java
@@ -48,8 +48,6 @@ public class QueryGroup extends AbstractDiffable<QueryGroup> implements ToXConte
     private final long updatedAtInMillis;
     private final Map<ResourceType, Object> resourceLimits;
 
-    // list of resources that are allowed to be present in the QueryGroup schema
-    public static final List<String> ALLOWED_RESOURCES = List.of("heap_allocations", "cpu");
 
     public QueryGroup(String name, QueryGroupMode resiliencyMode, Map<ResourceType, Object> resourceLimits) {
         this(name, UUIDs.randomBase64UUID(), resiliencyMode, resourceLimits, Instant.now().getMillis());
@@ -116,20 +114,13 @@ public void writeTo(StreamOutput out) throws IOException {
 
     private void validateResourceLimits(Map<ResourceType, Object> resourceLimits) {
         for (Map.Entry<ResourceType, Object> resource : resourceLimits.entrySet()) {
-            String resourceName = resource.getKey().getName();
             Double threshold = (Double) resource.getValue();
-            Objects.requireNonNull(resourceName, "resourceName can't be null");
-            Objects.requireNonNull(threshold, "resource limit threshold for" + resourceName + " : can't be null");
+            Objects.requireNonNull(resource.getKey(), "resourceName can't be null");
+            Objects.requireNonNull(threshold, "resource limit threshold for" + resource.getKey().getName() + " : can't be null");
 
             if (Double.compare(threshold, 1.0) > 0) {
                 throw new IllegalArgumentException("resource value should be less than 1.0");
             }
-
-            if (!ALLOWED_RESOURCES.contains(resourceName.toLowerCase(Locale.ROOT))) {
-                throw new IllegalArgumentException(
-                    "resource has to be valid, valid resources " + ALLOWED_RESOURCES.stream().reduce((x, e) -> x + ", " + e).get()
-                );
-            }
         }
     }
 
@@ -181,10 +172,8 @@ public static QueryGroup fromXContent(final XContentParser parser) throws IOExce
                     builder.mode(parser.text());
                 } else if (fieldName.equals("updatedAt")) {
                     builder.updatedAt(parser.longValue());
-                } else if (ALLOWED_RESOURCES.contains(fieldName)) {
-                    resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue());
                 } else {
-                    throw new IllegalArgumentException("unrecognised [field=" + fieldName + "] in QueryGroup");
+                    resourceLimits.put(ResourceType.fromName(fieldName), parser.doubleValue());
                 }
             }
         }

diff --git a/server/src/test/java/org/opensearch/cluster/ClusterModuleTests.java b/server/src/test/java/org/opensearch/cluster/ClusterModuleTests.java
@@ -33,6 +33,7 @@
 package org.opensearch.cluster;
 
 import org.opensearch.cluster.metadata.Metadata;
+import org.opensearch.cluster.metadata.QueryGroupMetadata;
 import org.opensearch.cluster.metadata.RepositoriesMetadata;
 import org.opensearch.cluster.routing.ShardRouting;
 import org.opensearch.cluster.routing.allocation.ExistingShardsAllocator;
@@ -69,6 +70,7 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.settings.SettingsModule;
 import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
 import org.opensearch.gateway.GatewayAllocator;
 import org.opensearch.plugins.ClusterPlugin;
 import org.opensearch.telemetry.metrics.noop.NoopMetricsRegistry;
@@ -327,6 +329,13 @@ public void testRejectsDuplicateExistingShardsAllocatorName() {
         );
     }
 
+    public void testQueryGroupMetadataRegister() {
+        List<NamedWriteableRegistry.Entry> customEntries = ClusterModule.getNamedWriteables();
+        assertTrue(customEntries.stream().anyMatch(entry ->
+            entry.categoryClass == Metadata.Custom.class && entry.name.equals(QueryGroupMetadata.TYPE)
+        ));
+    }
+
     private static ClusterPlugin existingShardsAllocatorPlugin(final String allocatorName) {
         return new ClusterPlugin() {
             @Override