DataNode replaceDataNodeWithNewOne(DataNode dataNode) {
      version += 1;
      DataNode newDataNode = testHardwareLayout.replaceDataNodeWithNewOne(dataNode);
      for (int i = 0; i < jsonPartitions.length(); i++) {
        JSONObject partitionInJson = jsonPartitions.getJSONObject(i);
        JSONArray replicasInJson = (JSONArray) partitionInJson.get("replicas");
        for (int j = 0; j < replicasInJson.length(); j++) {
          JSONObject replicaInJson = (JSONObject) replicasInJson.get(j);
          if (replicaInJson.getString("hostname").equals(dataNode.getHostname())
              && replicaInJson.getInt("port") == dataNode.getPort()) {
            replicasInJson.put(j,
                getJsonReplica(newDataNode.getHostname(), newDataNode.getPort(), replicaInJson.getString("mountPath")));
            break;
          }
        }
      }
      JSONObject jsonPartitionLayout =
          getJsonPartitionLayout(testHardwareLayout.getHardwareLayout().getClusterName(), version, partitionCount,
              jsonPartitions);
      this.partitionLayout =
          new PartitionLayout(testHardwareLayout.getHardwareLayout(), jsonPartitionLayout, clusterMapConfig);
      return newDataNode;
    }
Return the partition layout as a {@link JSONObject}
@param partitionClass the partition class that these partitions must belong to
@param partitionState the state of the newly created partitions
@return the partition layout as a {@link JSONObject}
@throws JSONException if there is an exception parsing or constructing JSON.
ambry-clustermap/src/test/java/com/github/ambry/clustermap/TestUtils.java