diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/indices.stats/30_segments.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/indices.stats/30_segments.yml
index 2d4d804063220..e4e820ff2b8bb 100644
--- a/rest-api-spec/src/main/resources/rest-api-spec/test/indices.stats/30_segments.yml
+++ b/rest-api-spec/src/main/resources/rest-api-spec/test/indices.stats/30_segments.yml
@@ -1,9 +1,5 @@
 ---
 setup:
-  - skip:
-      version: "all"
-      reason: "AwaitsFix: https://github.com/elastic/elasticsearch/issues/40331"
-
   - do:
       indices.create:
           index: test
@@ -47,6 +43,7 @@ setup:
   - do:
       indices.close:
         index: test
+        wait_for_active_shards: all
 
   - do:
       indices.stats: