From e365a255c7a23178b11b4d7f1084e0e1d8ab2c86 Mon Sep 17 00:00:00 2001
From: David Kyle <david.kyle@elastic.co>
Date: Tue, 20 Feb 2018 12:22:28 +0000
Subject: [PATCH] [ML] Filter non-ml attributes from data feed stats
 (elastic/x-pack-elasticsearch#3994)

Original commit: elastic/x-pack-elasticsearch@9cdc78c0005f1e6222737b3fac9690a78c0fc53c
---
 .../ml/action/GetDatafeedsStatsAction.java    |  4 +-
 .../GetDatafeedStatsActionResponseTests.java  | 54 +++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetDatafeedsStatsAction.java b/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetDatafeedsStatsAction.java
index b193ae28438..4219cb2a3ca 100644
--- a/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetDatafeedsStatsAction.java
+++ b/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetDatafeedsStatsAction.java
@@ -186,7 +186,9 @@ public class GetDatafeedsStatsAction extends Action<GetDatafeedsStatsAction.Requ
 
                     builder.startObject("attributes");
                     for (Map.Entry<String, String> entry : node.getAttributes().entrySet()) {
-                        builder.field(entry.getKey(), entry.getValue());
+                        if (entry.getKey().startsWith("ml.")) {
+                            builder.field(entry.getKey(), entry.getValue());
+                        }
                     }
                     builder.endObject();
                     builder.endObject();
diff --git a/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/action/GetDatafeedStatsActionResponseTests.java b/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/action/GetDatafeedStatsActionResponseTests.java
index 26716695c6e..0aaceeb28a7 100644
--- a/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/action/GetDatafeedStatsActionResponseTests.java
+++ b/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/action/GetDatafeedStatsActionResponseTests.java
@@ -7,16 +7,30 @@ package org.elasticsearch.xpack.core.ml.action;
 
 import org.elasticsearch.Version;
 import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.transport.TransportAddress;
+import org.elasticsearch.common.xcontent.ToXContent;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentHelper;
+import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.test.AbstractStreamableTestCase;
 import org.elasticsearch.xpack.core.ml.action.GetDatafeedsStatsAction.Response;
 import org.elasticsearch.xpack.core.ml.action.util.QueryPage;
 import org.elasticsearch.xpack.core.ml.datafeed.DatafeedConfig;
 import org.elasticsearch.xpack.core.ml.datafeed.DatafeedState;
 
+import java.io.IOException;
 import java.net.InetAddress;
 import java.util.ArrayList;
+import java.util.EnumSet;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.hasEntry;
+import static org.hamcrest.Matchers.hasKey;
+import static org.hamcrest.Matchers.is;
 
 public class GetDatafeedStatsActionResponseTests extends AbstractStreamableTestCase<Response> {
 
@@ -52,4 +66,44 @@ public class GetDatafeedStatsActionResponseTests extends AbstractStreamableTestC
         return new Response();
     }
 
+    @SuppressWarnings("unchecked")
+    public void testDatafeedStatsToXContent() throws IOException {
+        Map<String, String> attributes = new HashMap<>();
+        attributes.put("ml.enabled", "true");
+        attributes.put("ml.max_open_jobs", "5");
+        attributes.put("non-ml-attribute", "should be filtered out");
+        TransportAddress transportAddress = new TransportAddress(TransportAddress.META_ADDRESS, 9000);
+
+        DiscoveryNode node = new DiscoveryNode("df-node-name", "df-node-id", transportAddress, attributes,
+                EnumSet.noneOf(DiscoveryNode.Role.class),
+                Version.CURRENT);
+
+        Response.DatafeedStats stats = new Response.DatafeedStats("df-id", DatafeedState.STARTED, node, null);
+
+        XContentType xContentType = randomFrom(XContentType.values());
+        BytesReference bytes;
+        try (XContentBuilder builder = XContentBuilder.builder(xContentType.xContent())) {
+            stats.toXContent(builder, ToXContent.EMPTY_PARAMS);
+            bytes = builder.bytes();
+        }
+
+        Map<String, Object> dfStatsMap = XContentHelper.convertToMap(bytes, randomBoolean(), xContentType).v2();
+
+        assertThat(dfStatsMap.size(), is(equalTo(3)));
+        assertThat(dfStatsMap, hasEntry("datafeed_id", "df-id"));
+        assertThat(dfStatsMap, hasEntry("state", "started"));
+        assertThat(dfStatsMap, hasKey("node"));
+
+        Map<String, Object> nodeMap = (Map<String, Object>) dfStatsMap.get("node");
+        assertThat(nodeMap, hasEntry("id", "df-node-id"));
+        assertThat(nodeMap, hasEntry("name", "df-node-name"));
+        assertThat(nodeMap, hasKey("ephemeral_id"));
+        assertThat(nodeMap, hasKey("transport_address"));
+        assertThat(nodeMap, hasKey("attributes"));
+
+        Map<String, Object> nodeAttributes = (Map<String, Object>) nodeMap.get("attributes");
+        assertThat(nodeAttributes.size(), is(equalTo(2)));
+        assertThat(nodeAttributes, hasEntry("ml.enabled", "true"));
+        assertThat(nodeAttributes, hasEntry("ml.max_open_jobs", "5"));
+    }
 }