Enabled support for mBeans with non-string typed keys

jmxtrans · Feb 3, 2021 · 80f5805 · 80f5805
1 parent 711e04e
commit 80f5805
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 5 deletions.
diff --git a/src/main/java/org/jmxtrans/agent/Query.java b/src/main/java/org/jmxtrans/agent/Query.java
@@ -220,11 +220,12 @@ private void collectAndExportAttribute(MBeanServer mbeanServer, OutputWriter out
                     value = compositeData.get(key);
                 }
             } else if (attributeValue instanceof Map) {
-                Map<String, Object> mapData = (Map<String, Object>) attributeValue;
+                //The attribute key could be other types, e.g. Double/int
+                Map<Object, Object> mapData = (Map<Object, Object>) attributeValue;
                 if (key == null) {
                     // Get for all keys
-                    for (Map.Entry<String, Object> entry : mapData.entrySet()) {
-                        String key = entry.getKey();
+                    for (Map.Entry<Object, Object> entry : mapData.entrySet()) {
+                        String key = entry.getKey().toString();
                         value = entry.getValue();
                         processAttributeValue(outputWriter, objectName, attribute, key, value);
                     }

diff --git a/src/test/java/org/jmxtrans/agent/Mock.java b/src/test/java/org/jmxtrans/agent/Mock.java
@@ -23,6 +23,8 @@
  */
 package org.jmxtrans.agent;
 
+import java.util.HashMap;
+import java.util.Map;
 import javax.management.openmbean.*;
 import java.util.Arrays;
 import java.util.List;
@@ -89,4 +91,21 @@ public Integer[] getIntegerArray() {
     public List<Integer> getIntegerList() {
         return Arrays.asList(0, 1, 2, 3, 4, 5);
     }
+
+    @Override
+    public Map<String, Double> getStringMap() {
+        Map<String, Double> results = new HashMap<>();
+        results.put("foo", 0.01);
+        results.put("bar", 1.0);
+        return results;
+    }
+
+    @Override
+    public Map<Double, Double> getDoubleMap() {
+        Map<Double, Double> results = new HashMap<>();
+        results.put(0.01, 1.0);
+        results.put(0.02, 2.0);
+        results.put(0.03, 3.0);
+        return results;
+    }
 }
diff --git a/src/test/java/org/jmxtrans/agent/MockMBean.java b/src/test/java/org/jmxtrans/agent/MockMBean.java
@@ -23,6 +23,7 @@
  */
 package org.jmxtrans.agent;
 
+import java.util.Map;
 import javax.management.openmbean.CompositeData;
 import java.util.List;
 
@@ -50,4 +51,8 @@ public interface MockMBean {
     public Integer[] getIntegerArray();
 
     public List<Integer> getIntegerList();
+
+    public Map<String, Double> getStringMap();
+
+    public Map<Double, Double> getDoubleMap();
 }
diff --git a/src/test/java/org/jmxtrans/agent/QueryTest.java b/src/test/java/org/jmxtrans/agent/QueryTest.java
@@ -23,6 +23,7 @@
  */
 package org.jmxtrans.agent;
 
+import java.util.Map.Entry;
 import org.junit.*;
 
 import javax.annotation.Nonnull;
@@ -130,6 +131,34 @@ public void non_indexed_list_attribute_return_simple_result() throws Exception {
         }
     }
 
+    @Test
+    public void string_map_attribute_return_simple_result() throws Exception {
+        Query query = new Query("test:type=Mock,name=mock", "StringMap", null, null, null,
+            "StringMap_#key#", resultNameStrategy);
+        query.collectAndExport(mbeanServer, mockOutputWriter);
+
+        for (Entry<String, Double> e: mock.getStringMap().entrySet()) {
+            String name = "StringMap_" + e.getKey();
+            Object actual = mockOutputWriter.resultsByName.get(name);
+            assertThat("Result '" + name + "' is missing", actual, notNullValue());
+            assertThat("Result '" + name + "' type is invalid", actual, instanceOf(Number.class));
+        }
+    }
+
+    @Test
+    public void double_map_attribute_return_simple_result() throws Exception {
+        Query query = new Query("test:type=Mock,name=mock", "DoubleMap", null, null, null,
+            "DoubleMap_#key#", resultNameStrategy);
+        query.collectAndExport(mbeanServer, mockOutputWriter);
+
+        for (Entry<Double, Double> e: mock.getDoubleMap().entrySet()) {
+            String name = "DoubleMap_" + e.getKey().toString();
+            Object actual = mockOutputWriter.resultsByName.get(name);
+            assertThat("Result '" + name + "' is missing", actual, notNullValue());
+            assertThat("Result '" + name + "' type is invalid", actual, instanceOf(Number.class));
+        }
+    }
+
     @Test
     public void indexed_int_array_attribute_return_simple_result() throws Exception {
         Query query = new Query("test:type=Mock,name=mock", "IntArray", null, 1, null, "IntArray", resultNameStrategy);
@@ -189,7 +218,8 @@ public void query_objectname_with_null_attribute_returns_all_attributes() throws
         Query query = new Query("test:type=Mock,name=mock", null, null, resultNameStrategy);
         query.collectAndExport(mbeanServer, mockOutputWriter);
         Integer actualSize = mockOutputWriter.resultsByName.size();
-        assert (actualSize == 24);
+        int expectedSize = 24 + mock.getStringMap().size() + mock.getDoubleMap().size();
+        assert actualSize == expectedSize;
     }
 
     @Test
@@ -225,7 +255,7 @@ public MockOutputWriter(boolean failOnDuplicateResult) {
         @Override
         public void writeQueryResult(@Nonnull String name, @Nullable String type, @Nullable Object value) throws IOException {
             if (failOnDuplicateResult && resultsByName.containsKey(name)) {
-                fail("Result '" + name + "' already written");
+                fail("Result '" + name + "' with value '" + value.toString() + "' already written,");
             }
             resultsByName.put(name, value);
         }