From 30b20920b9c3419f8b9155a81b0623a8275bc143 Mon Sep 17 00:00:00 2001
From: Dimitris Athanasiou <dimitris@elastic.co>
Date: Fri, 5 Jul 2019 10:28:57 +0300
Subject: [PATCH] [7.x][ML] Report correct count for df-analytics get-stats API
 (#43969) (#43981)

The count should match the number of all df-analytics that
matched the id in the request. However, we set the count
to the number of df-analytics returned which was bound to the
`size` parameter. This commit fixes this by setting the count
to the count of the `get` response.
---
 ...TransportGetDataFrameAnalyticsStatsAction.java | 15 ++++++++++++---
 .../test/ml/data_frame_analytics_crud.yml         | 12 +++++++++---
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportGetDataFrameAnalyticsStatsAction.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportGetDataFrameAnalyticsStatsAction.java
index 8d139ba9b6c..884f741013b 100644
--- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportGetDataFrameAnalyticsStatsAction.java
+++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportGetDataFrameAnalyticsStatsAction.java
@@ -134,12 +134,21 @@ public class TransportGetDataFrameAnalyticsStatsAction
         LOGGER.debug("Get stats for data frame analytics [{}]", request.getId());
 
         ActionListener<GetDataFrameAnalyticsAction.Response> getResponseListener = ActionListener.wrap(
-            response -> {
-                List<String> expandedIds = response.getResources().results().stream().map(DataFrameAnalyticsConfig::getId)
+            getResponse -> {
+                List<String> expandedIds = getResponse.getResources().results().stream().map(DataFrameAnalyticsConfig::getId)
                     .collect(Collectors.toList());
                 request.setExpandedIds(expandedIds);
                 ActionListener<GetDataFrameAnalyticsStatsAction.Response> runningTasksStatsListener = ActionListener.wrap(
-                    runningTasksStatsResponse -> gatherStatsForStoppedTasks(request.getExpandedIds(), runningTasksStatsResponse, listener),
+                    runningTasksStatsResponse -> gatherStatsForStoppedTasks(request.getExpandedIds(), runningTasksStatsResponse,
+                        ActionListener.wrap(
+                            finalResponse -> {
+                                // While finalResponse has all the stats objects we need, we should report the count
+                                // from the get response
+                                QueryPage<Stats> finalStats = new QueryPage<>(finalResponse.getResponse().results(),
+                                    getResponse.getResources().count(), GetDataFrameAnalyticsAction.Response.RESULTS_FIELD);
+                                listener.onResponse(new GetDataFrameAnalyticsStatsAction.Response(finalStats));
+                            },
+                            listener::onFailure)),
                     listener::onFailure
                 );
                 super.doExecute(task, request, runningTasksStatsListener);
diff --git a/x-pack/plugin/src/test/resources/rest-api-spec/test/ml/data_frame_analytics_crud.yml b/x-pack/plugin/src/test/resources/rest-api-spec/test/ml/data_frame_analytics_crud.yml
index a171e879fa0..772c48e5474 100644
--- a/x-pack/plugin/src/test/resources/rest-api-spec/test/ml/data_frame_analytics_crud.yml
+++ b/x-pack/plugin/src/test/resources/rest-api-spec/test/ml/data_frame_analytics_crud.yml
@@ -776,6 +776,7 @@ setup:
       ml.get_data_frame_analytics_stats:
         id: "*"
   - match: { count: 3 }
+  - length: { data_frame_analytics : 3 }
   - match: { data_frame_analytics.0.id: "bar" }
   - match: { data_frame_analytics.0.state: "stopped" }
   - match: { data_frame_analytics.1.id: "foo-1" }
@@ -787,6 +788,7 @@ setup:
       ml.get_data_frame_analytics_stats:
         id: "foo-*"
   - match: { count: 2 }
+  - length: { data_frame_analytics : 2 }
   - match: { data_frame_analytics.0.id: "foo-1" }
   - match: { data_frame_analytics.0.state: "stopped" }
   - match: { data_frame_analytics.1.id: "foo-2" }
@@ -796,20 +798,23 @@ setup:
       ml.get_data_frame_analytics_stats:
         id: "bar"
   - match: { count: 1 }
+  - length: { data_frame_analytics : 1 }
   - match: { data_frame_analytics.0.id: "bar" }
   - match: { data_frame_analytics.0.state: "stopped" }
 
   - do:
       ml.get_data_frame_analytics_stats:
         from: 2
-  - match: { count: 1 }
+  - match: { count: 3 }
+  - length: { data_frame_analytics : 1 }
   - match: { data_frame_analytics.0.id: "foo-2" }
   - match: { data_frame_analytics.0.state: "stopped" }
 
   - do:
       ml.get_data_frame_analytics_stats:
         size: 2
-  - match: { count: 2 }
+  - match: { count: 3 }
+  - length: { data_frame_analytics : 2 }
   - match: { data_frame_analytics.0.id: "bar" }
   - match: { data_frame_analytics.0.state: "stopped" }
   - match: { data_frame_analytics.1.id: "foo-1" }
@@ -819,7 +824,8 @@ setup:
       ml.get_data_frame_analytics_stats:
         from: 1
         size: 1
-  - match: { count: 1 }
+  - match: { count: 3 }
+  - length: { data_frame_analytics : 1 }
   - match: { data_frame_analytics.0.id: "foo-1" }
   - match: { data_frame_analytics.0.state: "stopped" }