diff --git a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/diskbalancer/command/TestDiskBalancerCommand.java b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/diskbalancer/command/TestDiskBalancerCommand.java
index 1cdfcdcebc..d278983924 100644
--- a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/diskbalancer/command/TestDiskBalancerCommand.java
+++ b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/diskbalancer/command/TestDiskBalancerCommand.java
@@ -834,8 +834,13 @@ public void testDiskBalancerQueryWithoutSubmitAndMultipleNodes() throws Exceptio
           .getIpcPort(), dataNode2.getIpcPort());
       final String cmdLine = String.format("hdfs diskbalancer %s", queryArg);
       List<String> outputs = runCommand(cmdLine);
-      assertThat(outputs.get(1), containsString("localhost:" + dataNode1.getIpcPort()));
-      assertThat(outputs.get(6), containsString("localhost:" + dataNode2.getIpcPort()));
+      assertEquals(12,  outputs.size());
+      assertTrue("Expected outputs: " + outputs,
+          outputs.get(1).contains("localhost:" + dataNode1.getIpcPort()) ||
+              outputs.get(6).contains("localhost:" + dataNode1.getIpcPort()));
+      assertTrue("Expected outputs: " + outputs,
+          outputs.get(1).contains("localhost:" + dataNode2.getIpcPort()) ||
+              outputs.get(6).contains("localhost:" + dataNode2.getIpcPort()));
     } finally {
       miniDFSCluster.shutdown();
     }