diff --git a/test/py/tests/test_efi_capsule/test_capsule_firmware.py b/test/py/tests/test_efi_capsule/test_capsule_firmware.py
index 6e803f699f..a539134ec2 100644
--- a/test/py/tests/test_efi_capsule/test_capsule_firmware.py
+++ b/test/py/tests/test_efi_capsule/test_capsule_firmware.py
@@ -143,13 +143,14 @@ class TestEfiCapsuleFirmwareFit(object):
                 'fatls host 0:1 %s' % CAPSULE_INSTALL_DIR])
             assert 'Test01' in ''.join(output)
 
-        # reboot
-        u_boot_console.restart_uboot()
-
         capsule_early = u_boot_config.buildconfig.get(
             'config_efi_capsule_on_disk_early')
         capsule_auth = u_boot_config.buildconfig.get(
             'config_efi_capsule_authenticate')
+
+        # reboot
+        u_boot_console.restart_uboot(expect_earlyreset = capsule_early)
+
         with u_boot_console.log.section('Test Case 2-b, after reboot'):
             if not capsule_early:
                 # make sure that dfu_alt_info exists even persistent variables
@@ -162,7 +163,7 @@ class TestEfiCapsuleFirmwareFit(object):
 
                 # need to run uefi command to initiate capsule handling
                 output = u_boot_console.run_command(
-                    'env print -e Capsule0000')
+                    'env print -e Capsule0000', wait_for_reboot = True)
 
             output = u_boot_console.run_command_list([
                 'host bind 0 %s' % disk_img,
@@ -218,13 +219,14 @@ class TestEfiCapsuleFirmwareFit(object):
                 'fatls host 0:1 %s' % CAPSULE_INSTALL_DIR])
             assert 'Test02' in ''.join(output)
 
-        # reboot
-        u_boot_console.restart_uboot()
-
         capsule_early = u_boot_config.buildconfig.get(
             'config_efi_capsule_on_disk_early')
         capsule_auth = u_boot_config.buildconfig.get(
             'config_efi_capsule_authenticate')
+
+        # reboot
+        u_boot_console.restart_uboot(expect_earlyreset = capsule_early)
+
         with u_boot_console.log.section('Test Case 3-b, after reboot'):
             if not capsule_early:
                 # make sure that dfu_alt_info exists even persistent variables
@@ -237,7 +239,13 @@ class TestEfiCapsuleFirmwareFit(object):
 
                 # need to run uefi command to initiate capsule handling
                 output = u_boot_console.run_command(
-                    'env print -e Capsule0000')
+                    'env print -e Capsule0000', wait_for_reboot = True)
+
+            output = u_boot_console.run_command_list([
+                'env set dfu_alt_info "sf 0:0=u-boot-bin raw 0x100000 0x50000;u-boot-env raw 0x150000 0x200000"',
+                'host bind 0 %s' % disk_img,
+                'fatls host 0:1 %s' % CAPSULE_INSTALL_DIR])
+            assert 'Test02' not in ''.join(output)
 
             output = u_boot_console.run_command_list(['efidebug capsule esrt'])
 
@@ -293,13 +301,14 @@ class TestEfiCapsuleFirmwareFit(object):
                 'fatls host 0:1 %s' % CAPSULE_INSTALL_DIR])
             assert 'Test03' in ''.join(output)
 
-        # reboot
-        u_boot_console.restart_uboot()
-
         capsule_early = u_boot_config.buildconfig.get(
             'config_efi_capsule_on_disk_early')
         capsule_auth = u_boot_config.buildconfig.get(
             'config_efi_capsule_authenticate')
+
+        # reboot
+        u_boot_console.restart_uboot(expect_earlyreset = capsule_early)
+
         with u_boot_console.log.section('Test Case 4-b, after reboot'):
             if not capsule_early:
                 # make sure that dfu_alt_info exists even persistent variables
@@ -312,7 +321,13 @@ class TestEfiCapsuleFirmwareFit(object):
 
                 # need to run uefi command to initiate capsule handling
                 output = u_boot_console.run_command(
-                    'env print -e Capsule0000')
+                    'env print -e Capsule0000', wait_for_reboot = True)
+
+            output = u_boot_console.run_command_list([
+                'env set dfu_alt_info "sf 0:0=u-boot-bin raw 0x100000 0x50000;u-boot-env raw 0x150000 0x200000"',
+                'host bind 0 %s' % disk_img,
+                'fatls host 0:1 %s' % CAPSULE_INSTALL_DIR])
+            assert 'Test03' not in ''.join(output)
 
             output = u_boot_console.run_command_list(['efidebug capsule esrt'])
 
diff --git a/test/py/u_boot_console_base.py b/test/py/u_boot_console_base.py
index 384fd53c65..e84f4d74ef 100644
--- a/test/py/u_boot_console_base.py
+++ b/test/py/u_boot_console_base.py
@@ -140,7 +140,7 @@ class ConsoleBase(object):
         self.logstream.close()
 
     def run_command(self, cmd, wait_for_echo=True, send_nl=True,
-            wait_for_prompt=True):
+            wait_for_prompt=True, wait_for_reboot=False):
         """Execute a command via the U-Boot console.
 
         The command is always sent to U-Boot.
@@ -168,6 +168,8 @@ class ConsoleBase(object):
             wait_for_prompt: Boolean indicating whether to wait for the
                 command prompt to be sent by U-Boot. This typically occurs
                 immediately after the command has been executed.
+            wait_for_reboot: Boolean indication whether to wait for the
+                reboot U-Boot. If this is True, wait_for_prompt is ignored.
 
         Returns:
             If wait_for_prompt == False:
@@ -202,11 +204,48 @@ class ConsoleBase(object):
                                     self.bad_pattern_ids[m - 1])
             if not wait_for_prompt:
                 return
-            m = self.p.expect([self.prompt_compiled] + self.bad_patterns)
-            if m != 0:
-                self.at_prompt = False
-                raise Exception('Bad pattern found on console: ' +
-                                self.bad_pattern_ids[m - 1])
+            if wait_for_reboot:
+                bcfg = self.config.buildconfig
+                config_spl = bcfg.get('config_spl', 'n') == 'y'
+                config_spl_serial = bcfg.get('config_spl_serial',
+                                                 'n') == 'y'
+                env_spl_skipped = self.config.env.get('env__spl_skipped',
+                                                      False)
+                env_spl2_skipped = self.config.env.get('env__spl2_skipped',
+                                                       True)
+                if config_spl and config_spl_serial and not env_spl_skipped:
+                    m = self.p.expect([pattern_u_boot_spl_signon] +
+                                      self.bad_patterns)
+                    if m != 0:
+                        raise Exception('Bad pattern found on SPL console: ' +
+                                        self.bad_pattern_ids[m - 1])
+                if not env_spl2_skipped:
+                    m = self.p.expect([pattern_u_boot_spl2_signon] +
+                                      self.bad_patterns)
+                    if m != 0:
+                        raise Exception('Bad pattern found on SPL2 console: ' +
+                                        self.bad_pattern_ids[m - 1])
+                m = self.p.expect([pattern_u_boot_main_signon] + self.bad_patterns)
+                if m != 0:
+                    raise Exception('Bad pattern found on console: ' +
+                                    self.bad_pattern_ids[m - 1])
+                self.u_boot_version_string = self.p.after
+                while True:
+                    m = self.p.expect([self.prompt_compiled,
+                        pattern_stop_autoboot_prompt] + self.bad_patterns)
+                    if m == 0:
+                        break
+                    if m == 1:
+                        self.p.send(' ')
+                        continue
+                    raise Exception('Bad pattern found on console: ' +
+                                    self.bad_pattern_ids[m - 2])
+            else:
+                m = self.p.expect([self.prompt_compiled] + self.bad_patterns)
+                if m != 0:
+                    self.at_prompt = False
+                    raise Exception('Bad pattern found on console: ' +
+                                    self.bad_pattern_ids[m - 1])
             self.at_prompt = True
             self.at_prompt_logevt = self.logstream.logfile.cur_evt
             # Only strip \r\n; space/TAB might be significant if testing
@@ -321,7 +360,7 @@ class ConsoleBase(object):
         finally:
             self.p.timeout = orig_timeout
 
-    def ensure_spawned(self):
+    def ensure_spawned(self, expect_earlyreset=False):
         """Ensure a connection to a correctly running U-Boot instance.
 
         This may require spawning a new Sandbox process or resetting target
@@ -330,7 +369,8 @@ class ConsoleBase(object):
         This is an internal function and should not be called directly.
 
         Args:
-            None.
+            expect_earlyreset: This boot is expected to be reset in early
+                stage (before prompt). False by default.
 
         Returns:
             Nothing.
@@ -357,22 +397,29 @@ class ConsoleBase(object):
                                                   False)
             env_spl2_skipped = self.config.env.get('env__spl2_skipped',
                                                   True)
-            if config_spl and config_spl_serial and not env_spl_skipped:
-                m = self.p.expect([pattern_u_boot_spl_signon] +
-                                  self.bad_patterns)
+            if expect_earlyreset:
+                loop_num = 2
+            else:
+                loop_num = 1
+
+            while loop_num > 0:
+                loop_num -= 1
+                if config_spl and config_spl_serial and not env_spl_skipped:
+                    m = self.p.expect([pattern_u_boot_spl_signon] +
+                                      self.bad_patterns)
+                    if m != 0:
+                        raise Exception('Bad pattern found on SPL console: ' +
+                                        self.bad_pattern_ids[m - 1])
+                if not env_spl2_skipped:
+                    m = self.p.expect([pattern_u_boot_spl2_signon] +
+                                      self.bad_patterns)
+                    if m != 0:
+                        raise Exception('Bad pattern found on SPL2 console: ' +
+                                        self.bad_pattern_ids[m - 1])
+                m = self.p.expect([pattern_u_boot_main_signon] + self.bad_patterns)
                 if m != 0:
-                    raise Exception('Bad pattern found on SPL console: ' +
-                                    self.bad_pattern_ids[m - 1])
-            if not env_spl2_skipped:
-                m = self.p.expect([pattern_u_boot_spl2_signon] +
-                                  self.bad_patterns)
-                if m != 0:
-                    raise Exception('Bad pattern found on SPL2 console: ' +
+                    raise Exception('Bad pattern found on console: ' +
                                     self.bad_pattern_ids[m - 1])
-            m = self.p.expect([pattern_u_boot_main_signon] + self.bad_patterns)
-            if m != 0:
-                raise Exception('Bad pattern found on console: ' +
-                                self.bad_pattern_ids[m - 1])
             self.u_boot_version_string = self.p.after
             while True:
                 m = self.p.expect([self.prompt_compiled,
@@ -416,10 +463,10 @@ class ConsoleBase(object):
             pass
         self.p = None
 
-    def restart_uboot(self):
+    def restart_uboot(self, expect_earlyreset=False):
         """Shut down and restart U-Boot."""
         self.cleanup_spawn()
-        self.ensure_spawned()
+        self.ensure_spawned(expect_earlyreset)
 
     def get_spawn_output(self):
         """Return the start-up output from U-Boot
diff --git a/test/py/u_boot_console_sandbox.py b/test/py/u_boot_console_sandbox.py
index 7e1eb0e0b4..9cd9ccad30 100644
--- a/test/py/u_boot_console_sandbox.py
+++ b/test/py/u_boot_console_sandbox.py
@@ -57,11 +57,13 @@ class ConsoleSandbox(ConsoleBase):
         cmd += self.sandbox_flags
         return Spawn(cmd, cwd=self.config.source_dir)
 
-    def restart_uboot_with_flags(self, flags):
+    def restart_uboot_with_flags(self, flags, expect_earlyreset=False):
         """Run U-Boot with the given command-line flags
 
         Args:
             flags: List of flags to pass, each a string
+            expect_earlyreset: This boot is expected to be reset in early
+                stage (before prompt). False by default.
 
         Returns:
             A u_boot_spawn.Spawn object that is attached to U-Boot.
@@ -69,7 +71,7 @@ class ConsoleSandbox(ConsoleBase):
 
         try:
             self.sandbox_flags = flags
-            return self.restart_uboot()
+            return self.restart_uboot(expect_earlyreset)
         finally:
             self.sandbox_flags = []
 
