From 9e50dce46c22d78f28d3ea91dcdb6e9b3bcd6047 Mon Sep 17 00:00:00 2001
From: Anu Engineer <aengineer@apache.org>
Date: Fri, 1 Jun 2018 14:42:03 -0700
Subject: [PATCH] HDDS-144. Fix TestEndPoint#testHeartbeat and
 TestEndPoint#testRegister. Contributed by Shashikant Banerjee.

---
 .../ozone/container/common/ScmTestMock.java   | 32 +++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/hadoop-hdds/container-service/src/test/java/org/apache/hadoop/ozone/container/common/ScmTestMock.java b/hadoop-hdds/container-service/src/test/java/org/apache/hadoop/ozone/container/common/ScmTestMock.java
index 0ee6321c99..14da9601e4 100644
--- a/hadoop-hdds/container-service/src/test/java/org/apache/hadoop/ozone/container/common/ScmTestMock.java
+++ b/hadoop-hdds/container-service/src/test/java/org/apache/hadoop/ozone/container/common/ScmTestMock.java
@@ -16,6 +16,7 @@
  */
 package org.apache.hadoop.ozone.container.common;
 
+import com.google.common.base.Preconditions;
 import org.apache.hadoop.hdds.protocol.proto.StorageContainerDatanodeProtocolProtos;
 import org.apache.hadoop.hdds.scm.VersionInfo;
 import org.apache.hadoop.hdds.protocol.DatanodeDetails;
@@ -179,6 +180,7 @@ private void sleepIfNeeded() {
     List<SCMCommandProto>
         cmdResponses = new LinkedList<>();
     return SCMHeartbeatResponseProto.newBuilder().addAllCommands(cmdResponses)
+        .setDatanodeUUID(heartbeat.getDatanodeDetails().getUuid())
         .build();
   }
 
@@ -197,6 +199,7 @@ private void sleepIfNeeded() {
       throws IOException {
     rpcCount.incrementAndGet();
     updateNodeReport(datanodeDetailsProto, nodeReport);
+    updateContainerReport(containerReportsRequestProto, datanodeDetailsProto);
     sleepIfNeeded();
     return StorageContainerDatanodeProtocolProtos.SCMRegisteredResponseProto
         .newBuilder().setClusterID(UUID.randomUUID().toString())
@@ -227,6 +230,35 @@ public void updateNodeReport(DatanodeDetailsProto datanodeDetailsProto,
 
   }
 
+  /**
+   * Update the cotainerReport.
+   *
+   * @param reports Container report
+   * @param datanodeDetails DataNode Info
+   * @throws IOException
+   */
+  public void updateContainerReport(
+      StorageContainerDatanodeProtocolProtos.ContainerReportsProto reports,
+      DatanodeDetailsProto datanodeDetails) throws IOException {
+    Preconditions.checkNotNull(reports);
+    containerReportsCount.incrementAndGet();
+    DatanodeDetails datanode = DatanodeDetails.getFromProtoBuf(
+        datanodeDetails);
+    if (reports.getReportsCount() > 0) {
+      Map containers = nodeContainers.get(datanode);
+      if (containers == null) {
+        containers = new LinkedHashMap();
+        nodeContainers.put(datanode, containers);
+      }
+
+      for (StorageContainerDatanodeProtocolProtos.ContainerInfo report : reports
+          .getReportsList()) {
+        containers.put(report.getContainerID(), report);
+      }
+    }
+  }
+
+
   /**
    * Return the number of StorageReports of a datanode.
    * @param datanodeDetails