Enforce type safety for ExtensionNamedWriteableRegistry (#203)

* add generic type Signed-off-by: mloufra <mloufra@amazon.com> * move @SuppressWarnings at lowest scope possible Signed-off-by: mloufra <mloufra@amazon.com> * remove @SuppressWarnings Signed-off-by: mloufra <mloufra@amazon.com> Signed-off-by: mloufra <mloufra@amazon.com>
opensearch-project · Oct 26, 2022 · 35f8d2b · 35f8d2b
1 parent cf542d3
commit 35f8d2b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/src/main/java/org/opensearch/sdk/ExtensionNamedWriteableRegistry.java b/src/main/java/org/opensearch/sdk/ExtensionNamedWriteableRegistry.java
@@ -20,6 +20,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.common.io.stream.InputStreamStreamInput;
+import org.opensearch.common.io.stream.NamedWriteable;
 import org.opensearch.common.io.stream.NamedWriteableAwareStreamInput;
 import org.opensearch.common.io.stream.NamedWriteableRegistry;
 import org.opensearch.common.io.stream.NamedWriteableRegistryParseRequest;
@@ -79,12 +80,13 @@ private List<NamedWriteableRegistry.Entry> getNamedWriteables() {
      * @param request  The OpenSearch request to handle.
      * @return A response with a list of writeable names and fully qualified category class names to register within OpenSearch
      */
+    @SuppressWarnings("unchecked")
     public NamedWriteableRegistryResponse handleNamedWriteableRegistryRequest(OpenSearchRequest request) {
         logger.info("Registering Named Writeable Registry Request recieved from OpenSearch.");
         // Iterate through Extensions's named writeables and add to extension entries
-        Map<String, Class> extensionEntries = new HashMap<>();
+        Map<String, Class<? extends NamedWriteable>> extensionEntries = new HashMap<>();
         for (NamedWriteableRegistry.Entry entry : this.namedWriteables) {
-            extensionEntries.put(entry.name, entry.categoryClass);
+            extensionEntries.put(entry.name, (Class<? extends NamedWriteable>) entry.categoryClass);
         }
         NamedWriteableRegistryResponse namedWriteableRegistryResponse = new NamedWriteableRegistryResponse(extensionEntries);
         return namedWriteableRegistryResponse;
@@ -105,7 +107,7 @@ public ExtensionBooleanResponse handleNamedWriteableRegistryParseRequest(NamedWr
         boolean status = false;
 
         // Extract data from request and procress fully qualified category class name into class instance
-        Class categoryClass = request.getCategoryClass();
+        Class<? extends NamedWriteable> categoryClass = request.getCategoryClass();
         byte[] context = request.getContext();
 
         // Transform byte array context into an input stream

diff --git a/src/test/java/org/opensearch/sdk/TestNamedWriteableRegistryAPI.java b/src/test/java/org/opensearch/sdk/TestNamedWriteableRegistryAPI.java
@@ -135,12 +135,12 @@ public void testInvalidCategoryClass() throws UnknownHostException, IOException
         StreamInput in = new InputStreamStreamInput(input);
 
         // Category Class ExtensionRunner is not registered
-        NamedWriteableRegistryParseRequest request = new NamedWriteableRegistryParseRequest(ExtensionsRunner.class, in);
+        NamedWriteableRegistryParseRequest request = new NamedWriteableRegistryParseRequest(NamedWriteable.class, in);
         Exception e = expectThrows(
             Exception.class,
             () -> extensionNamedWriteableRegistry.handleNamedWriteableRegistryParseRequest(request)
         );
-        assertEquals(e.getMessage(), "Unknown NamedWriteable category [" + ExtensionsRunner.class.getName() + "]");
+        assertEquals(e.getMessage(), "Unknown NamedWriteable category [" + NamedWriteable.class.getName() + "]");
     }
 
     @Test