Test that shard write load is extracted

Author: nicktindall

server/src/test/java/org/elasticsearch/cluster/DiskUsageTests.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.cluster.routing.ShardRoutingHelper;
 import org.elasticsearch.cluster.routing.UnassignedInfo;
 import org.elasticsearch.index.Index;
+import org.elasticsearch.index.shard.IndexingStats;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.index.shard.ShardPath;
 import org.elasticsearch.index.store.StoreStats;
@@ -107,6 +108,7 @@ public void testFillShardLevelInfo() {
         Path test0Path = createTempDir().resolve("indices").resolve(index.getUUID()).resolve("0");
         CommonStats commonStats0 = new CommonStats();
         commonStats0.store = new StoreStats(100, 101, 0L);
+        commonStats0.indexing = randomIndexingStats();
         ShardRouting test_1 = ShardRouting.newUnassigned(
             new ShardId(index, 1),
             false,
@@ -119,8 +121,10 @@ public void testFillShardLevelInfo() {
         Path test1Path = createTempDir().resolve("indices").resolve(index.getUUID()).resolve("1");
         CommonStats commonStats1 = new CommonStats();
         commonStats1.store = new StoreStats(1000, 1001, 0L);
+        commonStats1.indexing = randomIndexingStats();
         CommonStats commonStats2 = new CommonStats();
         commonStats2.store = new StoreStats(1000, 999, 0L);
+        commonStats2.indexing = randomIndexingStats();
         ShardStats[] stats = new ShardStats[] {
             new ShardStats(test_0, new ShardPath(false, test0Path, test0Path, test_0.shardId()), commonStats0, null, null, null, false, 0),
             new ShardStats(test_1, new ShardPath(false, test1Path, test1Path, test_1.shardId()), commonStats1, null, null, null, false, 0),
@@ -166,6 +170,41 @@ public void testFillShardLevelInfo() {
                 hasEntry(ClusterInfo.NodeAndShard.from(test_1), test1Path.getParent().getParent().getParent().toAbsolutePath().toString())
             )
         );
+
+        assertThat(
+            shardWriteLoads,
+            equalTo(
+                Map.of(
+                    test_0.shardId(),
+                    commonStats0.indexing.getTotal().getPeakWriteLoad(),
+                    test_1.shardId(),
+                    Math.max(commonStats1.indexing.getTotal().getPeakWriteLoad(), commonStats2.indexing.getTotal().getPeakWriteLoad())
+                )
+            )
+        );
+    }
+
+    private IndexingStats randomIndexingStats() {
+        return new IndexingStats(
+            new IndexingStats.Stats(
+                randomNonNegativeLong(),
+                randomNonNegativeLong(),
+                randomNonNegativeLong(),
+                randomNonNegativeLong(),
+                randomNonNegativeLong(),
+                randomNonNegativeLong(),
+                randomMillisUpToYear9999(),
+                randomNonNegativeLong(),
+                randomNonNegativeLong(),
+                randomBoolean(),
+                randomNonNegativeLong(),
+                randomNonNegativeLong(),
+                randomNonNegativeLong(),
+                randomNonNegativeLong(),
+                randomDoubleBetween(0d, 10d, true),
+                randomDoubleBetween(0d, 10d, true)
+            )
+        );
     }
 
     public void testLeastAndMostAvailableDiskSpace() {