diff --git a/common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/RemoteBlockPushResolverSuite.java b/common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/RemoteBlockPushResolverSuite.java
index 0847121b0ccb0..e0b3315aad189 100644
--- a/common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/RemoteBlockPushResolverSuite.java
+++ b/common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/RemoteBlockPushResolverSuite.java
@@ -281,7 +281,7 @@ public void testDuplicateBlocksAreIgnoredWhenPrevStreamIsInProgress() throws IOE
     verifyMetrics(4, 0, 0, 0, 0, 0, 4);
   }
 
-  @Test(expected = RuntimeException.class)
+  @Test
   public void testFailureAfterData() throws IOException {
     StreamCallbackWithID stream =
       pushResolver.receiveBlockDataAsStream(
@@ -289,16 +289,13 @@ public void testFailureAfterData() throws IOException {
     stream.onData(stream.getID(), ByteBuffer.wrap(new byte[4]));
     stream.onFailure(stream.getID(), new RuntimeException("Forced Failure"));
     pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, NO_ATTEMPT_ID, 0, 0));
-    try {
-      pushResolver.getMergedBlockMeta(TEST_APP, 0, 0, 0);
-    } catch (RuntimeException e) {
-      assertTrue(e.getMessage().contains("is empty"));
-      verifyMetrics(4, 0, 0, 0, 0, 0, 4);
-      throw e;
-    }
+    RuntimeException e = assertThrows(RuntimeException.class,
+      () -> pushResolver.getMergedBlockMeta(TEST_APP, 0, 0, 0));
+    assertTrue(e.getMessage().contains("is empty"));
+    verifyMetrics(4, 0, 0, 0, 0, 0, 4);
   }
 
-  @Test(expected = RuntimeException.class)
+  @Test
   public void testFailureAfterMultipleDataBlocks() throws IOException {
     StreamCallbackWithID stream =
       pushResolver.receiveBlockDataAsStream(
@@ -308,13 +305,10 @@ public void testFailureAfterMultipleDataBlocks() throws IOException {
     stream.onData(stream.getID(), ByteBuffer.wrap(new byte[4]));
     stream.onFailure(stream.getID(), new RuntimeException("Forced Failure"));
     pushResolver.finalizeShuffleMerge(new FinalizeShuffleMerge(TEST_APP, NO_ATTEMPT_ID, 0, 0));
-    try {
-      pushResolver.getMergedBlockMeta(TEST_APP, 0, 0, 0);
-    } catch (RuntimeException e) {
-      assertTrue(e.getMessage().contains("is empty"));
-      verifyMetrics(9, 0, 0, 0, 0, 0, 9);
-      throw e;
-    }
+    RuntimeException e = assertThrows(RuntimeException.class,
+      () -> pushResolver.getMergedBlockMeta(TEST_APP, 0, 0, 0));
+    assertTrue(e.getMessage().contains("is empty"));
+    verifyMetrics(9, 0, 0, 0, 0, 0, 9);
   }
 
   @Test
diff --git a/dev/checkstyle.xml b/dev/checkstyle.xml
index 343eaa4cfda7c..5af15318081a6 100644
--- a/dev/checkstyle.xml
+++ b/dev/checkstyle.xml
@@ -180,6 +180,10 @@
               value="Avoid using com.google.common.io.Files.createTempDir() due to CVE-2020-8908.
                 Use org.apache.spark.network.util.JavaUtils.createTempDir() instead." />
         </module>
+        <module name="RegexpSinglelineJava">
+            <property name="format" value="@Test\(expected"/>
+            <property name="message" value="Please use the `assertThrows` method to test for exceptions." />
+        </module>
         <module name="IllegalImport">
             <property name="illegalPkgs" value="org.apache.log4j" />
         </module>