Fix TestDiskBalancerWithMockMover.

Change-Id: Id31b53479d96b569256ddf03c5ea4f5fdc017bef

hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/diskbalancer/TestDiskBalancerWithMockMover.java

@@ -22,6 +22,7 @@ package org.apache.hadoop.hdfs.server.diskbalancer;
 import org.apache.hadoop.thirdparty.com.google.common.base.Preconditions;
 import java.util.function.Supplier;
 import org.apache.commons.codec.digest.DigestUtils;
+import org.junit.jupiter.api.Assertions;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.apache.hadoop.conf.Configuration;
@@ -46,7 +47,6 @@ import org.apache.hadoop.util.Time;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
-import org.junit.rules.ExpectedException;
 
 import java.io.IOException;
 import java.net.URI;
@@ -88,11 +88,11 @@ public class TestDiskBalancerWithMockMover {
         .setMover(blockMover)
         .build();
 
-    thrown.expect(DiskBalancerException.class);
-    thrown.expect(new DiskBalancerResultVerifier(DiskBalancerException
-        .Result.DISK_BALANCER_NOT_ENABLED));
-
-    balancer.queryWorkStatus();
+    final DiskBalancerException thrown =
+        Assertions.assertThrows(DiskBalancerException.class,
+            balancer::queryWorkStatus);
+    Assertions.assertEquals(thrown.getResult(),
+        DiskBalancerException.Result.DISK_BALANCER_NOT_ENABLED);
   }
 
   /**
@@ -143,10 +143,12 @@ public class TestDiskBalancerWithMockMover {
     // ask block mover to get stuck in copy block
     mockMoverHelper.getBlockMover().setSleep();
     executeSubmitPlan(plan, balancer);
-    thrown.expect(DiskBalancerException.class);
-    thrown.expect(new DiskBalancerResultVerifier(DiskBalancerException
-        .Result.PLAN_ALREADY_IN_PROGRESS));
-    executeSubmitPlan(plan, balancer);
+    final DiskBalancerException thrown =
+        Assertions.assertThrows(DiskBalancerException.class, () -> {
+          executeSubmitPlan(plan, balancer);
+        });
+    Assertions.assertEquals(thrown.getResult(),
+        DiskBalancerException.Result.PLAN_ALREADY_IN_PROGRESS);
 
     // Not needed but this is the cleanup step.
     mockMoverHelper.getBlockMover().clearSleep();
@@ -183,10 +185,12 @@ public class TestDiskBalancerWithMockMover {
     DiskBalancer balancer = mockMoverHelper.getBalancer();
 
     plan.setTimeStamp(Time.now() - (32 * millisecondInAnHour));
-    thrown.expect(DiskBalancerException.class);
-    thrown.expect(new DiskBalancerResultVerifier(DiskBalancerException
-        .Result.OLD_PLAN_SUBMITTED));
-    executeSubmitPlan(plan, balancer);
+    final DiskBalancerException thrown =
+        Assertions.assertThrows(DiskBalancerException.class, () -> {
+          executeSubmitPlan(plan, balancer);
+        });
+    Assertions.assertEquals(thrown.getResult(),
+        DiskBalancerException.Result.OLD_PLAN_SUBMITTED);
   }
 
   @Test
@@ -195,12 +199,13 @@ public class TestDiskBalancerWithMockMover {
     NodePlan plan = mockMoverHelper.getPlan();
     DiskBalancer balancer = mockMoverHelper.getBalancer();
 
-    thrown.expect(DiskBalancerException.class);
-    thrown.expect(new DiskBalancerResultVerifier(DiskBalancerException
-        .Result.INVALID_PLAN_VERSION));
-
-    // Plan version is invalid -- there is no version 0.
-    executeSubmitPlan(plan, balancer, 0);
+    final DiskBalancerException thrown =
+        Assertions.assertThrows(DiskBalancerException.class, () -> {
+          // Plan version is invalid -- there is no version 0.
+          executeSubmitPlan(plan, balancer, 0);
+        });
+    Assertions.assertEquals(thrown.getResult(),
+        DiskBalancerException.Result.INVALID_PLAN_VERSION);
   }
 
   @Test
@@ -211,11 +216,12 @@ public class TestDiskBalancerWithMockMover {
     String planJson = plan.toJson();
     String planID = DigestUtils.sha1Hex(planJson);
 
-    thrown.expect(DiskBalancerException.class);
-    thrown.expect(new DiskBalancerResultVerifier(DiskBalancerException
-        .Result.INVALID_PLAN));
-
-    balancer.submitPlan(planID, 1, "no-plan-file.json", null, false);
+    final DiskBalancerException thrown =
+        Assertions.assertThrows(DiskBalancerException.class, () -> {
+          balancer.submitPlan(planID, 1, "no-plan-file.json", null, false);
+        });
+    Assertions.assertEquals(thrown.getResult(),
+        DiskBalancerException.Result.INVALID_PLAN);
   }
 
   @Test
@@ -227,15 +233,16 @@ public class TestDiskBalancerWithMockMover {
 
     String planJson = plan.toJson();
     String planID = DigestUtils.sha1Hex(planJson);
-    char repChar = planID.charAt(0);
-    repChar++;
-
-    thrown.expect(DiskBalancerException.class);
-    thrown.expect(new DiskBalancerResultVerifier(DiskBalancerException
-        .Result.INVALID_PLAN_HASH));
-    balancer.submitPlan(planID.replace(planID.charAt(0), repChar),
-        1, PLAN_FILE, planJson, false);
-
+    char repCharValid = planID.charAt(0);
+    final char repChar = ++repCharValid;
+
+    final DiskBalancerException thrown =
+        Assertions.assertThrows(DiskBalancerException.class, () -> {
+          balancer.submitPlan(planID.replace(planID.charAt(0), repChar),
+              1, PLAN_FILE, planJson, false);
+        });
+    Assertions.assertEquals(thrown.getResult(),
+        DiskBalancerException.Result.INVALID_PLAN_HASH);
   }
 
   /**
@@ -267,12 +274,14 @@ public class TestDiskBalancerWithMockMover {
     executeSubmitPlan(plan, balancer);
 
     // Send a Wrong cancellation request.
-    char first = planID.charAt(0);
-    first++;
-    thrown.expect(DiskBalancerException.class);
-    thrown.expect(new DiskBalancerResultVerifier(DiskBalancerException
-        .Result.NO_SUCH_PLAN));
-    balancer.cancelPlan(planID.replace(planID.charAt(0), first));
+    char firstValid = planID.charAt(0);
+    final char first = ++firstValid;
+    final DiskBalancerException thrown =
+        Assertions.assertThrows(DiskBalancerException.class, () -> {
+          balancer.cancelPlan(planID.replace(planID.charAt(0), first));
+        });
+    Assertions.assertEquals(thrown.getResult(),
+        DiskBalancerException.Result.NO_SUCH_PLAN);
 
     // Now cancel the real one
     balancer.cancelPlan(planID);