Covers OffHeap.

Author: viirya

core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java

@@ -33,6 +33,7 @@
 
 import org.apache.spark.SparkConf;
 import org.apache.spark.executor.ShuffleWriteMetrics;
+import org.apache.spark.memory.MemoryMode;
 import org.apache.spark.memory.TestMemoryConsumer;
 import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.memory.TestMemoryManager;
@@ -671,7 +672,8 @@ public void testPeakMemoryUsed() {
   @Test
   public void avoidDeadlock() throws InterruptedException {
     memoryManager.limit(PAGE_SIZE_BYTES);
-    TestMemoryConsumer c1 = new TestMemoryConsumer(taskMemoryManager);
+    MemoryMode mode = useOffHeapMemoryAllocator() ? MemoryMode.OFF_HEAP: MemoryMode.ON_HEAP;
+    TestMemoryConsumer c1 = new TestMemoryConsumer(taskMemoryManager, mode);
     BytesToBytesMap map =
       new BytesToBytesMap(taskMemoryManager, blockManager, serializerManager, 1, 0.5, 1024);