diff --git a/tests/vb2_ui_tests.c b/tests/vb2_ui_tests.c
index 6f5ca68..325e842 100644
--- a/tests/vb2_ui_tests.c
+++ b/tests/vb2_ui_tests.c
@@ -15,7 +15,7 @@
 #include "test_common.h"
 #include "vboot_kernel.h"
 
-/* Fixed value for ignoring some checks. */
+/* Fixed value for ignoring some checks */
 #define MOCK_IGNORE 0xffffu
 
 /* Mock data */
@@ -41,9 +41,12 @@
 
 static int mock_calls_until_shutdown;
 
+/* Iteration counter starts from 0
+   Mock inputs should response according to this */
+static int mock_iters;
+
 static uint32_t mock_key[64];
 static int mock_key_trusted[64];
-static int mock_key_count;
 static int mock_key_total;
 
 static uint32_t mock_get_timer_last;
@@ -61,7 +64,6 @@
 
 static vb2_error_t mock_vbtlk_retval[32];
 static uint32_t mock_vbtlk_expected_flag[32];
-static int mock_vbtlk_count;
 static int mock_vbtlk_total;
 
 static void add_mock_key(uint32_t press, int trusted)
@@ -180,13 +182,13 @@
 	else
 		mock_calls_until_shutdown = 10;
 
+	/* For iteration counter */
+	mock_iters = -1;  /* Accumulates at the beginning of iterations */
+
 	/* For VbExKeyboardRead */
 	memset(mock_key, 0, sizeof(mock_key));
 	memset(mock_key_trusted, 0, sizeof(mock_key_trusted));
-	mock_key_count = 0;
 	mock_key_total = 0;
-	/* Avoid iteration #0 which has a screen change by global action */
-	add_mock_keypress(0);
 
 	/* For vboot_audio.h */
 	mock_get_timer_last = 0;
@@ -206,8 +208,15 @@
 	/* For VbTryLoadKernel */
 	memset(mock_vbtlk_retval, 0, sizeof(mock_vbtlk_retval));
 	memset(mock_vbtlk_expected_flag, 0, sizeof(mock_vbtlk_expected_flag));
-	mock_vbtlk_count = 0;
 	mock_vbtlk_total = 0;
+
+	/* Avoid Iteration #0 */
+	add_mock_keypress(0);
+	if (t == FOR_MANUAL_RECOVERY)
+		add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND,
+			       VB_DISK_FLAG_REMOVABLE);
+	else
+		add_mock_vbtlk(VB2_ERROR_MOCK, 0);
 }
 
 /* Mock functions */
@@ -260,14 +269,15 @@
 
 uint32_t VbExKeyboardReadWithFlags(uint32_t *key_flags)
 {
-	if (mock_key_count < mock_key_total) {
+	mock_iters++;
+	if (mock_iters < mock_key_total) {
 		if (key_flags != NULL) {
-			if (mock_key_trusted[mock_key_count])
+			if (mock_key_trusted[mock_iters])
 				*key_flags = VB_KEY_FLAG_TRUSTED_KEYBOARD;
 			else
 				*key_flags = 0;
 		}
-		return mock_key[mock_key_count++];
+		return mock_key[mock_iters];
 	}
 
 	return 0;
@@ -319,19 +329,16 @@
 
 vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t get_info_flags)
 {
-	if (mock_vbtlk_total == 0) {
-		TEST_TRUE(0, "  VbTryLoadKernel is not allowed!");
-		return VB2_ERROR_MOCK;
-	}
+	int i = mock_iters;
 
 	/* Return last entry if called too many times */
-	if (mock_vbtlk_count >= mock_vbtlk_total)
-		mock_vbtlk_count = mock_vbtlk_total - 1;
+	if (i >= mock_vbtlk_total)
+		i = mock_vbtlk_total - 1;
 
-	TEST_EQ(mock_vbtlk_expected_flag[mock_vbtlk_count], get_info_flags,
+	TEST_EQ(mock_vbtlk_expected_flag[i], get_info_flags,
 		"  unexpected get_info_flags");
 
-	return mock_vbtlk_retval[mock_vbtlk_count++];
+	return mock_vbtlk_retval[i];
 }
 
 /* Tests */
@@ -350,7 +357,7 @@
 	TEST_TRUE(mock_get_timer_last - mock_time_start >=
 		  30 * VB2_MSEC_PER_SEC, "  finished delay");
 	TEST_EQ(mock_vbexbeep_called, 2, "  beeped twice");
-	TEST_EQ(mock_vbtlk_count, mock_vbtlk_total, "  used up mock_vbtlk");
+	TEST_TRUE(mock_iters >= mock_vbtlk_total, "  used up mock_vbtlk");
 
 	/* Proceed to USB after timeout */
 	reset_common_data(FOR_DEVELOPER);
@@ -365,7 +372,7 @@
 	TEST_TRUE(mock_get_timer_last - mock_time_start >=
 		  30 * VB2_MSEC_PER_SEC, "  finished delay");
 	TEST_EQ(mock_vbexbeep_called, 2, "  beeped twice");
-	TEST_EQ(mock_vbtlk_count, mock_vbtlk_total, "  used up mock_vbtlk");
+	TEST_TRUE(mock_iters >= mock_vbtlk_total, "  used up mock_vbtlk");
 
 	/* Default boot USB not allowed, don't boot */
 	reset_common_data(FOR_DEVELOPER);
@@ -378,7 +385,7 @@
 	TEST_TRUE(mock_get_timer_last - mock_time_start >=
 		  30 * VB2_MSEC_PER_SEC, "  finished delay");
 	TEST_EQ(mock_vbexbeep_called, 2, "  beeped twice");
-	TEST_EQ(mock_vbtlk_count, mock_vbtlk_total, "  used up mock_vbtlk");
+	TEST_TRUE(mock_iters >= mock_vbtlk_total, "  used up mock_vbtlk");
 
 	VB2_DEBUG("...done.\n");
 }
@@ -424,7 +431,6 @@
 
 	/* Timeout, shutdown */
 	reset_common_data(FOR_MANUAL_RECOVERY);
-	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
 	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
 		"timeout, shutdown");
 	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
@@ -435,8 +441,6 @@
 	if (!DETACHABLE) {
 		reset_common_data(FOR_MANUAL_RECOVERY);
 		add_mock_keypress(VB_BUTTON_POWER_SHORT_PRESS);
-		add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND,
-			       VB_DISK_FLAG_REMOVABLE);
 		TEST_EQ(vb2_manual_recovery_menu(ctx),
 			VB2_REQUEST_SHUTDOWN,
 			"power button short pressed = shutdown");
@@ -448,7 +452,6 @@
 	/* Item 1 = phone recovery */
 	reset_common_data(FOR_MANUAL_RECOVERY);
 	add_mock_keypress(VB_KEY_ENTER);
-	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
 	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
 		"phone recovery");
 	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
@@ -461,7 +464,6 @@
 	reset_common_data(FOR_MANUAL_RECOVERY);
 	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
-	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
 	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
 		"external disk recovery");
 	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,