From 1ff7043a1fbf16108371594c0b0c5902d6f50b6b Mon Sep 17 00:00:00 2001
From: Sanford Ryza <sandy@apache.org>
Date: Thu, 1 Aug 2013 06:10:48 +0000
Subject: [PATCH] YARN-758. Augment MockNM to use multiple cores (Karthik
 Kambatla via Sandy Ryza)

git-svn-id: https://svn.apache.org/repos/asf/hadoop/common/trunk@1509086 13f79535-47bb-0310-9956-ffa450edef68
---
 hadoop-yarn-project/CHANGES.txt                     |  3 +++
 .../hadoop/yarn/server/resourcemanager/MockNM.java  | 13 ++++++++++++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/hadoop-yarn-project/CHANGES.txt b/hadoop-yarn-project/CHANGES.txt
index c3821daec6c..47466b872c7 100644
--- a/hadoop-yarn-project/CHANGES.txt
+++ b/hadoop-yarn-project/CHANGES.txt
@@ -27,6 +27,9 @@ Release 2.3.0 - UNRELEASED
 
   BUG FIXES
 
+    YARN-758. Augment MockNM to use multiple cores (Karthik Kambatla via
+    Sandy Ryza)
+
 Release 2.1.1-beta - UNRELEASED
 
   INCOMPATIBLE CHANGES
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockNM.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockNM.java
index 0aea4b692be..1258bb86707 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockNM.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockNM.java
@@ -30,6 +30,7 @@ import org.apache.hadoop.yarn.api.records.ContainerState;
 import org.apache.hadoop.yarn.api.records.ContainerStatus;
 import org.apache.hadoop.yarn.api.records.NodeId;
 import org.apache.hadoop.yarn.api.records.Resource;
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.server.api.protocolrecords.NodeHeartbeatRequest;
 import org.apache.hadoop.yarn.server.api.protocolrecords.NodeHeartbeatResponse;
 import org.apache.hadoop.yarn.server.api.protocolrecords.RegisterNodeManagerRequest;
@@ -45,14 +46,24 @@ public class MockNM {
   private int responseId;
   private NodeId nodeId;
   private final int memory;
-  private final int vCores = 1;
+  private final int vCores;
   private ResourceTrackerService resourceTracker;
   private final int httpPort = 2;
   private MasterKey currentContainerTokenMasterKey;
   private MasterKey currentNMTokenMasterKey;
 
   public MockNM(String nodeIdStr, int memory, ResourceTrackerService resourceTracker) {
+    // scale vcores based on the requested memory
+    this(nodeIdStr, memory,
+        Math.max(1, (memory * YarnConfiguration.DEFAULT_NM_VCORES) /
+            YarnConfiguration.DEFAULT_NM_PMEM_MB),
+        resourceTracker);
+  }
+
+  public MockNM(String nodeIdStr, int memory, int vcores,
+                ResourceTrackerService resourceTracker) {
     this.memory = memory;
+    this.vCores = vcores;
     this.resourceTracker = resourceTracker;
     String[] splits = nodeIdStr.split(":");
     nodeId = BuilderUtils.newNodeId(splits[0], Integer.parseInt(splits[1]));