fix tests

core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java

@@ -66,6 +66,7 @@ public class UnsafeShuffleWriterSuite {
   TaskMemoryManager taskMemoryManager;
   final HashPartitioner hashPartitioner = new HashPartitioner(NUM_PARTITITONS);
   File mergedOutputFile;
+  File indexFile;
   File mapStatusFile;
   File tempDir;
   long[] partitionSizesInMergedFile;
@@ -106,9 +107,11 @@ public void setUp() throws IOException {
     MockitoAnnotations.initMocks(this);
     tempDir = Utils.createTempDir("test", "test");
     mergedOutputFile = File.createTempFile("mergedoutput", "", tempDir);
+    indexFile = File.createTempFile("shuffle",".index", tempDir);
     mapStatusFile = File.createTempFile("shuffle", ".mapstatus", tempDir);
     // the ShuffleOutputCoordinator requires that this file does not exist
     mergedOutputFile.delete();
+    indexFile.delete();
     mapStatusFile.delete();
     partitionSizesInMergedFile = null;
     tmpShuffleFilesCreated.clear();
@@ -171,6 +174,7 @@ public OutputStream answer(InvocationOnMock invocation) throws Throwable {
     when(blockManager.shuffleServerId()).thenReturn(BlockManagerId$.MODULE$.apply("1", "a.b.c", 1));
 
     when(shuffleBlockResolver.getDataFile(anyInt(), anyInt())).thenReturn(mergedOutputFile);
+    when(shuffleBlockResolver.getIndexFile(anyInt(), anyInt())).thenReturn(indexFile);
     doAnswer(new Answer<File>() {
       @Override
       public File answer(InvocationOnMock invocationOnMock) throws Throwable {
@@ -191,18 +195,6 @@ public Tuple2<TempShuffleBlockId, File> answer(
         }
       });
 
-    when(diskBlockManager.getFile(any(BlockId.class))).thenAnswer(
-      new Answer<File>() {
-        @Override
-        public File answer(InvocationOnMock invocationOnMock) throws Throwable {
-          File f = File.createTempFile("shuffleFile",".index", tempDir);
-          // the ShuffleOutputCoordinator requires that this file does not exist
-          f.delete();
-          return f;
-        }
-      }
-    );
-
     when(taskContext.taskMetrics()).thenReturn(taskMetrics);
     when(taskContext.internalMetricsToAccumulators()).thenReturn(null);
 

core/src/test/scala/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriterSuite.scala

@@ -76,6 +76,7 @@ class BypassMergeSortShuffleWriterSuite extends SparkFunSuite with BeforeAndAfte
     when(dependency.serializer).thenReturn(Some(new JavaSerializer(conf)))
     when(taskContext.taskMetrics()).thenReturn(taskMetrics)
     when(blockResolver.getDataFile(0, 0)).thenReturn(outputFile)
+    when(blockResolver.getIndexFile(0, 0)).thenReturn(indexFile)
     // the index file will be empty, but that is fine for these tests
     when(blockResolver.writeIndexFile(anyInt(), anyInt(), any())).thenAnswer(new Answer[File] {
       override def answer(invocationOnMock: InvocationOnMock): File = {
@@ -120,12 +121,7 @@ class BypassMergeSortShuffleWriterSuite extends SparkFunSuite with BeforeAndAfte
     when(diskBlockManager.getFile(any[BlockId])).thenAnswer(
       new Answer[File] {
         override def answer(invocation: InvocationOnMock): File = {
-          val blk = invocation.getArguments.head.asInstanceOf[BlockId]
-          if (blk == new ShuffleIndexBlockId(0, 0, IndexShuffleBlockResolver.NOOP_REDUCE_ID)) {
-            indexFile
-          } else {
-            blockIdToFileMap.get(invocation.getArguments.head.asInstanceOf[BlockId]).get
-          }
+          blockIdToFileMap.get(invocation.getArguments.head.asInstanceOf[BlockId]).get
         }
     })
   }