diff --git a/hadoop-hdfs-project/hadoop-hdfs/CHANGES_HDFS-5535.txt b/hadoop-hdfs-project/hadoop-hdfs/CHANGES_HDFS-5535.txt
index f7c5ab4aa71..2e3e1c3016c 100644
--- a/hadoop-hdfs-project/hadoop-hdfs/CHANGES_HDFS-5535.txt
+++ b/hadoop-hdfs-project/hadoop-hdfs/CHANGES_HDFS-5535.txt
@@ -99,3 +99,8 @@ HDFS-5535 subtasks:
     HDFS-5498. Improve datanode startup time. (kihwal)
 
     HDFS-6000. Avoid saving namespace when starting rolling upgrade. (jing9)
+
+    HDFS-6017. Query the status of rolling upgrade in the preparation stage in
+    TestRollingUpgrade and TestRollingUpgradeRollback. (Haohui Mai via
+    Arpit Agarwal)
+
diff --git a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestRollingUpgrade.java b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestRollingUpgrade.java
index 1efbe827fd6..6dfbdc08ecd 100644
--- a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestRollingUpgrade.java
+++ b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestRollingUpgrade.java
@@ -17,6 +17,9 @@
  */
 package org.apache.hadoop.hdfs;
 
+import java.io.File;
+import java.io.IOException;
+
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.conf.Configuration;
@@ -30,15 +33,11 @@ import org.apache.hadoop.hdfs.qjournal.MiniJournalCluster;
 import org.apache.hadoop.hdfs.qjournal.MiniQJMHACluster;
 import org.apache.hadoop.hdfs.server.common.HdfsServerConstants.StartupOption;
 import org.apache.hadoop.hdfs.server.datanode.DataNode;
-import org.apache.hadoop.hdfs.server.namenode.NNStorage;
+import org.apache.hadoop.hdfs.server.namenode.FSImage;
 import org.apache.hadoop.hdfs.tools.DFSAdmin;
 import org.junit.Assert;
 import org.junit.Test;
 
-import java.io.File;
-import java.io.FilenameFilter;
-import java.io.IOException;
-
 
 /**
  * This class tests rolling upgrade.
@@ -384,32 +383,6 @@ public class TestRollingUpgrade {
     }
   }
 
-  public static boolean existRollbackFsImage(NNStorage storage)
-      throws IOException {
-    final FilenameFilter filter = new FilenameFilter() {
-      @Override
-      public boolean accept(File dir, String name) {
-        return name.indexOf(NNStorage.NameNodeFile.IMAGE_ROLLBACK.getName()) != -1;
-      }
-    };
-    final int total = 10;
-    int retry = 0;
-    while (retry++ < total) {
-      for (int i = 0; i < storage.getNumStorageDirs(); i++) {
-        File dir = storage.getStorageDir(i).getCurrentDir();
-        int l = dir.list(filter).length;
-        if (l > 0) {
-          return true;
-        }
-      }
-      try {
-        Thread.sleep(1000);
-      } catch (InterruptedException e) {
-      }
-    }
-    return false;
-  }
-
   @Test (timeout = 300000)
   public void testFinalize() throws Exception {
     final Configuration conf = new HdfsConfiguration();
@@ -431,26 +404,42 @@ public class TestRollingUpgrade {
       DistributedFileSystem dfs = dfsCluster.getFileSystem(0);
       dfs.mkdirs(foo);
 
-      NNStorage storage = dfsCluster.getNamesystem(0).getFSImage()
-          .getStorage();
+      FSImage fsimage = dfsCluster.getNamesystem(0).getFSImage();
 
       // start rolling upgrade
       RollingUpgradeInfo info = dfs.rollingUpgrade(RollingUpgradeAction.PREPARE);
       Assert.assertTrue(info.isStarted());
       dfs.mkdirs(bar);
+
+      queryForPreparation(dfs);
+
       // The NN should have a copy of the fsimage in case of rollbacks.
-      Assert.assertTrue(existRollbackFsImage(storage));
+      Assert.assertTrue(fsimage.hasRollbackFSImage());
 
       info = dfs.rollingUpgrade(RollingUpgradeAction.FINALIZE);
       Assert.assertTrue(info.isFinalized());
       Assert.assertTrue(dfs.exists(foo));
 
       // Once finalized, there should be no more fsimage for rollbacks.
-      Assert.assertFalse(existRollbackFsImage(storage));
+      Assert.assertFalse(fsimage.hasRollbackFSImage());
     } finally {
       if (cluster != null) {
         cluster.shutdown();
       }
     }
   }
+
+  static void queryForPreparation(DistributedFileSystem dfs) throws IOException,
+      InterruptedException {
+    RollingUpgradeInfo info;
+    int retries = 0;
+    while (retries < 10) {
+      info = dfs.rollingUpgrade(RollingUpgradeAction.QUERY);
+      if (info.createdRollbackImages()) {
+        break;
+      }
+      Thread.sleep(1000);
+      ++retries;
+    }
+  }
 }
diff --git a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestRollingUpgradeRollback.java b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestRollingUpgradeRollback.java
index 93eb3744fc2..c2e9d7cff2a 100644
--- a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestRollingUpgradeRollback.java
+++ b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestRollingUpgradeRollback.java
@@ -228,10 +228,14 @@ public class TestRollingUpgradeRollback {
       dfs.mkdirs(bar);
       dfs.close();
 
-      NNStorage storage0 = dfsCluster.getNameNode(0).getFSImage().getStorage();
-      NNStorage storage1 = dfsCluster.getNameNode(1).getFSImage().getStorage();
-      Assert.assertTrue(TestRollingUpgrade.existRollbackFsImage(storage0));
-      Assert.assertTrue(TestRollingUpgrade.existRollbackFsImage(storage1));
+      TestRollingUpgrade.queryForPreparation(dfs);
+
+      // If the query returns true, both active and the standby NN should have
+      // rollback fsimage ready.
+      Assert.assertTrue(dfsCluster.getNameNode(0).getFSImage()
+          .hasRollbackFSImage());
+      Assert.assertTrue(dfsCluster.getNameNode(1).getFSImage()
+          .hasRollbackFSImage());
       
       // rollback NN0
       dfsCluster.restartNameNode(0, true, "-rollingUpgrade",