diff --git a/common/network-yarn/src/test/java/org/apache/spark/network/yarn/YarnShuffleServiceSuite.java b/common/network-yarn/src/test/java/org/apache/spark/network/yarn/YarnShuffleServiceSuite.java
deleted file mode 100644
index 09bc4d840367..000000000000
--- a/common/network-yarn/src/test/java/org/apache/spark/network/yarn/YarnShuffleServiceSuite.java
+++ /dev/null
@@ -1,61 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.network.yarn;
-
-import org.junit.Test;
-
-import static org.junit.Assert.*;
-import static org.mockito.Mockito.*;
-
-import org.apache.spark.network.shuffle.ExternalBlockHandler;
-import org.apache.spark.network.shuffle.MergedShuffleFileManager;
-import org.apache.spark.network.shuffle.RemoteBlockPushResolver;
-import org.apache.spark.network.util.TransportConf;
-
-public class YarnShuffleServiceSuite {
-
-  @Test
-  public void testCreateDefaultMergedShuffleFileManagerInstance() {
-    TransportConf mockConf = mock(TransportConf.class);
-    when(mockConf.mergedShuffleFileManagerImpl()).thenReturn(
-      "org.apache.spark.network.shuffle.ExternalBlockHandler$NoOpMergedShuffleFileManager");
-    MergedShuffleFileManager mergeMgr = YarnShuffleService.newMergedShuffleFileManagerInstance(
-      mockConf);
-    assertTrue(mergeMgr instanceof ExternalBlockHandler.NoOpMergedShuffleFileManager);
-  }
-
-  @Test
-  public void testCreateRemoteBlockPushResolverInstance() {
-    TransportConf mockConf = mock(TransportConf.class);
-    when(mockConf.mergedShuffleFileManagerImpl()).thenReturn(
-      "org.apache.spark.network.shuffle.RemoteBlockPushResolver");
-    MergedShuffleFileManager mergeMgr = YarnShuffleService.newMergedShuffleFileManagerInstance(
-      mockConf);
-    assertTrue(mergeMgr instanceof RemoteBlockPushResolver);
-  }
-
-  @Test
-  public void testInvalidClassNameOfMergeManagerWillUseNoOpInstance() {
-    TransportConf mockConf = mock(TransportConf.class);
-    when(mockConf.mergedShuffleFileManagerImpl()).thenReturn(
-      "org.apache.spark.network.shuffle.NotExistent");
-    MergedShuffleFileManager mergeMgr = YarnShuffleService.newMergedShuffleFileManagerInstance(
-      mockConf);
-    assertTrue(mergeMgr instanceof ExternalBlockHandler.NoOpMergedShuffleFileManager);
-  }
-}
diff --git a/resource-managers/yarn/src/test/scala/org/apache/spark/network/yarn/YarnShuffleServiceSuite.scala b/resource-managers/yarn/src/test/scala/org/apache/spark/network/yarn/YarnShuffleServiceSuite.scala
index a6a302ad5df9..c2bdd971a0fe 100644
--- a/resource-managers/yarn/src/test/scala/org/apache/spark/network/yarn/YarnShuffleServiceSuite.scala
+++ b/resource-managers/yarn/src/test/scala/org/apache/spark/network/yarn/YarnShuffleServiceSuite.scala
@@ -34,6 +34,7 @@ import org.apache.hadoop.service.ServiceStateException
 import org.apache.hadoop.yarn.api.records.ApplicationId
 import org.apache.hadoop.yarn.conf.YarnConfiguration
 import org.apache.hadoop.yarn.server.api.{ApplicationInitializationContext, ApplicationTerminationContext}
+import org.mockito.Mockito.{mock, when}
 import org.scalatest.BeforeAndAfterEach
 import org.scalatest.concurrent.Eventually._
 import org.scalatest.matchers.must.Matchers
@@ -42,8 +43,9 @@ import org.scalatest.matchers.should.Matchers._
 import org.apache.spark.SecurityManager
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.internal.config._
-import org.apache.spark.network.shuffle.ShuffleTestAccessor
+import org.apache.spark.network.shuffle.{ExternalBlockHandler, RemoteBlockPushResolver, ShuffleTestAccessor}
 import org.apache.spark.network.shuffle.protocol.ExecutorShuffleInfo
+import org.apache.spark.network.util.TransportConf
 import org.apache.spark.util.Utils
 
 class YarnShuffleServiceSuite extends SparkFunSuite with Matchers with BeforeAndAfterEach {
@@ -411,4 +413,27 @@ class YarnShuffleServiceSuite extends SparkFunSuite with Matchers with BeforeAnd
     ))
   }
 
+  test("create default merged shuffle file manager instance") {
+    val mockConf = mock(classOf[TransportConf])
+    when(mockConf.mergedShuffleFileManagerImpl).thenReturn(
+      "org.apache.spark.network.shuffle.ExternalBlockHandler$NoOpMergedShuffleFileManager")
+    val mergeMgr = YarnShuffleService.newMergedShuffleFileManagerInstance(mockConf)
+    assert(mergeMgr.isInstanceOf[ExternalBlockHandler.NoOpMergedShuffleFileManager])
+  }
+
+  test("create remote block push resolver instance") {
+    val mockConf = mock(classOf[TransportConf])
+    when(mockConf.mergedShuffleFileManagerImpl).thenReturn(
+      "org.apache.spark.network.shuffle.RemoteBlockPushResolver")
+    val mergeMgr = YarnShuffleService.newMergedShuffleFileManagerInstance(mockConf)
+    assert(mergeMgr.isInstanceOf[RemoteBlockPushResolver])
+  }
+
+  test("invalid class name of merge manager will use noop instance") {
+    val mockConf = mock(classOf[TransportConf])
+    when(mockConf.mergedShuffleFileManagerImpl).thenReturn(
+      "org.apache.spark.network.shuffle.NotExistent")
+    val mergeMgr = YarnShuffleService.newMergedShuffleFileManagerInstance(mockConf)
+    assert(mergeMgr.isInstanceOf[ExternalBlockHandler.NoOpMergedShuffleFileManager])
+  }
 }