From bae6e3c66d5f10f7f0e20a63a20670596b3e90bb Mon Sep 17 00:00:00 2001
From: Niels Bauman <33722607+nielsbauman@users.noreply.github.com>
Date: Mon, 23 Jun 2025 12:22:48 -0300
Subject: [PATCH] Fix data stream stats YAML test (#129813)

Occasional shard allocation issues were causing the YAML tests to fail
because the shard that had the document with the max timestamp in it
would be unavailable.

Fixes #118217
---
 .../test/data_stream/120_data_streams_stats.yml       | 11 +++++++++++
 muted-tests.yml                                       |  3 ---
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/modules/data-streams/src/yamlRestTest/resources/rest-api-spec/test/data_stream/120_data_streams_stats.yml b/modules/data-streams/src/yamlRestTest/resources/rest-api-spec/test/data_stream/120_data_streams_stats.yml
index 7ee9e656007e..e60741c63af2 100644
--- a/modules/data-streams/src/yamlRestTest/resources/rest-api-spec/test/data_stream/120_data_streams_stats.yml
+++ b/modules/data-streams/src/yamlRestTest/resources/rest-api-spec/test/data_stream/120_data_streams_stats.yml
@@ -185,8 +185,19 @@ setup:
         body:  { "@timestamp": 1593639468350 }
         refresh: true
 
+  # Wait for the shards to be allocated to avoid missing results in the stats request
+  # See https://github.com/elastic/elasticsearch/issues/118217#issuecomment-2996315057 for more context
+  - do:
+      cluster.health:
+        index: simple-data-stream1,simple-data-stream2
+        wait_for_status: green
+        wait_for_no_relocating_shards: true
+
   - do:
       indices.data_streams_stats: {}
+  - match: { _shards.total: 3 }
+  - match: { _shards.successful: 3 }
+  - match: { _shards.failed: 0 }
   - match: { data_stream_count: 2 }
   - match: { backing_indices: 3 }
   - length: { data_streams: 2 }
diff --git a/muted-tests.yml b/muted-tests.yml
index 2b59966cd5bc..00317f504578 100644
--- a/muted-tests.yml
+++ b/muted-tests.yml
@@ -102,9 +102,6 @@ tests:
 - class: org.elasticsearch.packaging.test.ArchiveTests
   method: test51AutoConfigurationWithPasswordProtectedKeystore
   issue: https://github.com/elastic/elasticsearch/issues/118212
-- class: org.elasticsearch.datastreams.DataStreamsClientYamlTestSuiteIT
-  method: test {p0=data_stream/120_data_streams_stats/Multiple data stream}
-  issue: https://github.com/elastic/elasticsearch/issues/118217
 - class: org.elasticsearch.xpack.ccr.rest.ShardChangesRestIT
   method: testShardChangesNoOperation
   issue: https://github.com/elastic/elasticsearch/issues/118800