Fix tests

Author: nicktindall

server/src/main/java/org/elasticsearch/cluster/ClusterInfo.java

@@ -322,7 +322,8 @@ public String toString() {
 
     // exposed for tests, computed here rather than exposing all the collections separately
     int getChunkCount() {
-        return leastAvailableSpaceUsage.size() + shardSizes.size() + shardDataSetSizes.size() + dataPath.size() + reservedSpace.size() + 6;
+        return leastAvailableSpaceUsage.size() + shardSizes.size() + shardDataSetSizes.size() + dataPath.size() + reservedSpace.size()
+            + nodesHeapUsage.size() + 7;
     }
 
     public record NodeAndShard(String nodeId, ShardId shardId) implements Writeable {

server/src/test/java/org/elasticsearch/action/admin/cluster/allocation/DesiredBalanceResponseTests.java

@@ -337,7 +337,7 @@ public void testToXContent() throws IOException {
         Map<String, Object> clusterInfo = (Map<String, Object>) json.get("cluster_info");
         assertThat(
             clusterInfo.keySet(),
-            containsInAnyOrder("nodes", "shard_paths", "shard_sizes", "shard_data_set_sizes", "reserved_sizes")
+            containsInAnyOrder("nodes", "shard_paths", "shard_sizes", "shard_data_set_sizes", "reserved_sizes", "heap_usage")
         );
     }