diff --git a/hadoop-hdfs-project/hadoop-hdfs/CHANGES_HDFS-5535.txt b/hadoop-hdfs-project/hadoop-hdfs/CHANGES_HDFS-5535.txt
index a9dc4a564f..e021cd9947 100644
--- a/hadoop-hdfs-project/hadoop-hdfs/CHANGES_HDFS-5535.txt
+++ b/hadoop-hdfs-project/hadoop-hdfs/CHANGES_HDFS-5535.txt
@@ -114,3 +114,6 @@ HDFS-5535 subtasks:
 
     HDFS-6024. Test whether the NN will reject the downgrade if it has a
     fsimage from a newer release. (Haohui Mai via jing9)
+
+    HDFS-6026. Fix TestDFSUpgrade and TestDataNodeRollingUpgrade.
+    (jing9 via szetszwo)
diff --git a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestDFSUpgrade.java b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestDFSUpgrade.java
index 73006baed1..ee9e91dc73 100644
--- a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestDFSUpgrade.java
+++ b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestDFSUpgrade.java
@@ -235,7 +235,7 @@ public void testUpgrade() throws Exception {
       // make sure that rolling upgrade cannot be started
       try {
         final DistributedFileSystem dfs = cluster.getFileSystem();
-        dfs.setSafeMode(SafeModeAction.SAFEMODE_LEAVE);
+        dfs.setSafeMode(SafeModeAction.SAFEMODE_ENTER);
         dfs.rollingUpgrade(RollingUpgradeAction.PREPARE);
         fail();
       } catch(RemoteException re) {
@@ -378,7 +378,7 @@ public void testUpgrade() throws Exception {
       // make sure that rolling upgrade cannot be started
       try {
         final DistributedFileSystem dfs = cluster.getFileSystem();
-        dfs.setSafeMode(SafeModeAction.SAFEMODE_LEAVE);
+        dfs.setSafeMode(SafeModeAction.SAFEMODE_ENTER);
         dfs.rollingUpgrade(RollingUpgradeAction.PREPARE);
         fail();
       } catch(RemoteException re) {
diff --git a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/datanode/TestDataNodeRollingUpgrade.java b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/datanode/TestDataNodeRollingUpgrade.java
index d548da4e10..5299c7c2ec 100644
--- a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/datanode/TestDataNodeRollingUpgrade.java
+++ b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/datanode/TestDataNodeRollingUpgrade.java
@@ -41,6 +41,7 @@
 import org.apache.hadoop.hdfs.protocol.BlockLocalPathInfo;
 import org.apache.hadoop.hdfs.protocol.ExtendedBlock;
 import org.apache.hadoop.hdfs.protocol.LocatedBlocks;
+import org.apache.hadoop.hdfs.protocol.HdfsConstants.SafeModeAction;
 import org.apache.hadoop.hdfs.server.namenode.NameNode;
 import org.apache.hadoop.hdfs.tools.DFSAdmin;
 import org.junit.Test;
@@ -194,6 +195,7 @@ public void testDatanodeRollingUpgradeWithFinalize() throws Exception {
       DFSTestUtil.createFile(fs, testFile1, FILE_SIZE, REPL_FACTOR, SEED);
       DFSTestUtil.createFile(fs, testFile2, FILE_SIZE, REPL_FACTOR, SEED);
 
+      fs.setSafeMode(SafeModeAction.SAFEMODE_ENTER);
       startRollingUpgrade();
       File blockFile = getBlockForFile(testFile2, true);
       File trashFile = getTrashFileForBlock(blockFile, false);
@@ -220,6 +222,7 @@ public void testDatanodeRollingUpgradeWithRollback() throws Exception {
       DFSTestUtil.createFile(fs, testFile1, FILE_SIZE, REPL_FACTOR, SEED);
       String fileContents1 = DFSTestUtil.readFile(fs, testFile1);
 
+      fs.setSafeMode(SafeModeAction.SAFEMODE_ENTER);
       startRollingUpgrade();
 
       File blockFile = getBlockForFile(testFile1, true);