vboot/ui: Refine mock functions in UI tests

Rearrange and split some action tests.

Rewrite mock_ui_context related functions, and allow change_screen to
change to real screens in vb2_ui_action_tests.c.

BUG=b:146399181
TEST=make clean && make runtests
TEST=make clean && DETACHABLE=1; make runtests
BRANCH=none

Signed-off-by: Hsuan Ting Chen <roccochen@chromium.org>
Change-Id: I1692c9fcdf19073f8a9810dfab769df4c87224a9
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2209780
Reviewed-by: Yu-Ping Wu <yupingso@chromium.org>

3 files changed, 189 insertions, 175 deletions

diff --git a/tests/vb2_ui_action_tests.c b/tests/vb2_ui_action_tests.c
index f438f01f..166830c4 100644
--- a/tests/vb2_ui_action_tests.c
+++ b/tests/vb2_ui_action_tests.c
@@ -18,14 +18,13 @@
 #define MOCK_IGNORE 0xffffu
 
 /* Mock screen index for testing screen utility functions. */
-#define MOCK_NO_SCREEN 0xef0
-#define MOCK_SCREEN_BASE 0xeff
-#define MOCK_SCREEN_MENU 0xfff
-#define MOCK_SCREEN_TARGET0 0xff0
-#define MOCK_SCREEN_TARGET1 0xff1
-#define MOCK_SCREEN_TARGET2 0xff2
-#define MOCK_SCREEN_TARGET3 0xff3
-#define MOCK_SCREEN_TARGET4 0xff4
+#define MOCK_NO_SCREEN 0xef00
+#define MOCK_SCREEN_BASE 0xef10
+#define MOCK_SCREEN_MENU 0xef11
+#define MOCK_SCREEN_TARGET0 0xef20
+#define MOCK_SCREEN_TARGET1 0xef21
+#define MOCK_SCREEN_TARGET2 0xef22
+#define MOCK_SCREEN_TARGET3 0xef23
 
 /* Mock data */
 struct display_call {
@@ -54,26 +53,41 @@ static int mock_key_trusted[64];
 static int mock_key_count;
 static int mock_key_total;
 
-static enum vb2_screen mock_get_screen_info_last;
+static int mock_get_screen_info_called;
 
 static vb2_error_t mock_vbtlk_retval;
 static uint32_t mock_vbtlk_expected_flag;
 
+/* Mock actions */
+static uint32_t mock_action_called;
+static vb2_error_t mock_action_countdown(struct vb2_ui_context *ui)
+{
+	if (++mock_action_called >= 10)
+		return VB2_SUCCESS;
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
+static vb2_error_t mock_action_change_screen(struct vb2_ui_context *ui)
+{
+	return vb2_ui_change_screen(ui, MOCK_SCREEN_BASE);
+}
+
 /* Mock screens */
+struct vb2_screen_info mock_screen_temp;
 const struct vb2_menu_item mock_empty_menu[] = {};
-struct vb2_screen_info mock_screen_blank = {
+const struct vb2_screen_info mock_screen_blank = {
 	.id = VB2_SCREEN_BLANK,
 	.name = "mock_screen_blank",
 	.num_items = ARRAY_SIZE(mock_empty_menu),
 	.items = mock_empty_menu,
 };
-struct vb2_screen_info mock_screen_base = {
+const struct vb2_screen_info mock_screen_base = {
 	.id = MOCK_SCREEN_BASE,
 	.name = "mock_screen_base: menuless screen",
 	.num_items = ARRAY_SIZE(mock_empty_menu),
 	.items = mock_empty_menu,
 };
-struct vb2_menu_item mock_screen_menu_items[] = {
+const struct vb2_menu_item mock_screen_menu_items[] = {
 	{
 		.text = "option 0",
 		.target = MOCK_SCREEN_TARGET0,
@@ -124,27 +138,6 @@ const struct vb2_screen_info mock_screen_target3 = {
 	.num_items = ARRAY_SIZE(mock_empty_menu),
 	.items = mock_empty_menu,
 };
-const struct vb2_screen_info mock_screen_target4 = {
-	.id = MOCK_SCREEN_TARGET4,
-	.name = "mock_screen_target4",
-	.num_items = ARRAY_SIZE(mock_empty_menu),
-	.items = mock_empty_menu,
-};
-
-/* Mock actions */
-static uint32_t global_action_called;
-static vb2_error_t global_action_countdown(struct vb2_ui_context *ui)
-{
-	if (++global_action_called >= 10)
-		return VB2_SUCCESS;
-	return VB2_REQUEST_UI_CONTINUE;
-}
-
-static vb2_error_t global_action_change_screen(struct vb2_ui_context *ui)
-{
-	vb2_ui_change_screen(ui, MOCK_SCREEN_BASE);
-	return VB2_REQUEST_UI_CONTINUE;
-}
 
 static void screen_state_eq(const struct vb2_screen_state *state,
 			    enum vb2_screen screen,
@@ -249,19 +242,17 @@ static void reset_common_data(void)
 	/* For check_shutdown_request */
 	mock_calls_until_shutdown = 10;
 
-	/* Mock ui_context based on mock screens */
-	mock_ui_context = (struct vb2_ui_context){
-		.ctx = ctx,
-		.root_screen = &mock_screen_blank,
-		.state = (struct vb2_screen_state){
-			.screen = &mock_screen_blank,
-			.selected_item = 0,
-			.disabled_item_mask = 0,
-		},
-		.locale_id = 0,
-		.key = 0,
-
+	/* Reset mock_screen_temp for test by test temporary screen_info */
+	mock_screen_temp = (struct vb2_screen_info){
+	      .id = MOCK_NO_SCREEN,
+	      .name = "mock_screen_temp",
+	      .num_items = ARRAY_SIZE(mock_empty_menu),
+	      .items = mock_empty_menu,
 	};
+
+	/* Mock ui_context based on mock screens */
+	memset(&mock_ui_context, 0, sizeof(mock_ui_context));
+	mock_ui_context.ctx = ctx;
 	mock_state = &mock_ui_context.state;
 
 	/* For vb2ex_display_ui */
@@ -275,11 +266,11 @@ static void reset_common_data(void)
 	mock_key_count = 0;
 	mock_key_total = 0;
 
-	/* For global actions */
-	global_action_called = 0;
+	/* For mock actions */
+	mock_action_called = 0;
 
-	/* For vb2_get_screen_info */
-	mock_get_screen_info_last = -1;
+	/* For chagen_screen and vb2_get_screen_info */
+	mock_get_screen_info_called = 0;
 
 	/* For VbTryLoadKernel */
 	mock_vbtlk_retval = VB2_ERROR_MOCK;
@@ -305,7 +296,7 @@ uint32_t VbExIsShutdownRequested(void)
 
 const struct vb2_screen_info *vb2_get_screen_info(enum vb2_screen screen)
 {
-	mock_get_screen_info_last = screen;
+	mock_get_screen_info_called++;
 
 	switch ((int)screen) {
 	case VB2_SCREEN_BLANK:
@@ -322,12 +313,11 @@ const struct vb2_screen_info *vb2_get_screen_info(enum vb2_screen screen)
 		return &mock_screen_target2;
 	case MOCK_SCREEN_TARGET3:
 		return &mock_screen_target3;
-	case MOCK_SCREEN_TARGET4:
-		return &mock_screen_target4;
 	case MOCK_NO_SCREEN:
 		return NULL;
 	default:
-		return &mock_screen_blank;
+		mock_screen_temp.id = screen;
+		return &mock_screen_temp;
 	}
 }
 
@@ -381,54 +371,50 @@ vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t get_info_flags)
 {
 	TEST_EQ(mock_vbtlk_expected_flag, get_info_flags,
 		"  unexpected get_info_flags");
-
 	return mock_vbtlk_retval;
 }
 
 /* Tests */
-static void menu_action_tests(void)
+static void menu_up_action_tests(void)
 {
-	int i, target_id;
-	char test_name[256];
-
-	VB2_DEBUG("Testing menu actions...\n");
+	VB2_DEBUG("Testing menu_up_action...\n");
 
-	/* Valid menu_up_action */
+	/* Valid action */
 	reset_common_data();
 	mock_state->screen = &mock_screen_menu;
 	mock_state->selected_item = 2;
 	mock_ui_context.key = VB_KEY_UP;
 	TEST_EQ(menu_up_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"valid menu_up_action");
+		"valid action");
 	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 1, MOCK_IGNORE);
 
-	/* Valid menu_up_action with mask */
+	/* Valid action with mask */
 	reset_common_data();
 	mock_state->screen = &mock_screen_menu;
 	mock_state->selected_item = 2;
 	mock_state->disabled_item_mask = 0x0a;  /* 0b01010 */
 	mock_ui_context.key = VB_KEY_UP;
 	TEST_EQ(menu_up_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"valid menu_up_action with mask");
+		"valid action with mask");
 	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 0, MOCK_IGNORE);
 
-	/* Invalid menu_up_action (blocked) */
+	/* Invalid action (blocked) */
 	reset_common_data();
 	mock_state->screen = &mock_screen_menu;
 	mock_state->selected_item = 0;
 	mock_ui_context.key = VB_KEY_UP;
 	TEST_EQ(menu_up_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"invalid menu_up_action (blocked)");
+		"invalid action (blocked)");
 	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 0, MOCK_IGNORE);
 
-	/* Invalid menu_up_action (blocked by mask) */
+	/* Invalid action (blocked by mask) */
 	reset_common_data();
 	mock_state->screen = &mock_screen_menu;
 	mock_state->selected_item = 2;
 	mock_state->disabled_item_mask = 0x0b;  /* 0b01011 */
 	mock_ui_context.key = VB_KEY_UP;
 	TEST_EQ(menu_up_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"invalid menu_up_action (blocked by mask)");
+		"invalid action (blocked by mask)");
 	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
 
 	/* Ignore volume-up when not DETACHABLE */
@@ -443,42 +429,49 @@ static void menu_action_tests(void)
 		screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
 	}
 
-	/* Valid menu_down_action */
+	VB2_DEBUG("...done.\n");
+}
+
+static void menu_down_action_tests(void)
+{
+	VB2_DEBUG("Testing menu_down_action...\n");
+
+	/* Valid action */
 	reset_common_data();
 	mock_state->screen = &mock_screen_menu;
 	mock_state->selected_item = 2;
 	mock_ui_context.key = VB_KEY_DOWN;
 	TEST_EQ(menu_down_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"valid menu_down_action");
+		"valid action");
 	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 3, MOCK_IGNORE);
 
-	/* Valid menu_down_action with mask */
+	/* Valid action with mask */
 	reset_common_data();
 	mock_state->screen = &mock_screen_menu;
 	mock_state->selected_item = 2;
 	mock_state->disabled_item_mask = 0x0a;  /* 0b01010 */
 	mock_ui_context.key = VB_KEY_DOWN;
 	TEST_EQ(menu_down_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"valid menu_down_action with mask");
+		"valid action with mask");
 	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 4, MOCK_IGNORE);
 
-	/* Invalid menu_down_action (blocked) */
+	/* Invalid action (blocked) */
 	reset_common_data();
 	mock_state->screen = &mock_screen_menu;
 	mock_state->selected_item = 4;
 	mock_ui_context.key = VB_KEY_DOWN;
 	TEST_EQ(menu_down_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"invalid menu_down_action (blocked)");
+		"invalid action (blocked)");
 	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 4, MOCK_IGNORE);
 
-	/* Invalid menu_down_action (blocked by mask) */
+	/* Invalid action (blocked by mask) */
 	reset_common_data();
 	mock_state->screen = &mock_screen_menu;
 	mock_state->selected_item = 2;
 	mock_state->disabled_item_mask = 0x1a;  /* 0b11010 */
 	mock_ui_context.key = VB_KEY_DOWN;
 	TEST_EQ(menu_down_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"invalid menu_down_action (blocked by mask)");
+		"invalid action (blocked by mask)");
 	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
 
 	/* Ignore volume-down when not DETACHABLE */
@@ -493,7 +486,17 @@ static void menu_action_tests(void)
 		screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
 	}
 
-	/* menu_select_action with no item screen */
+	VB2_DEBUG("...done.\n");
+}
+
+static void menu_select_action_tests(void)
+{
+	int i, target_id;
+	char test_name[256];
+
+	VB2_DEBUG("Testing menu_select_action...\n");
+
+	/* select action with no item screen */
 	reset_common_data();
 	mock_state->screen = &mock_screen_base;
 	mock_ui_context.key = VB_KEY_ENTER;
@@ -515,7 +518,7 @@ static void menu_action_tests(void)
 		screen_state_eq(mock_state, target_id, 0, MOCK_IGNORE);
 	}
 
-	/* Try to select no target item */
+	/* Try to select no target item (target 4) */
 	reset_common_data();
 	mock_state->screen = &mock_screen_menu;
 	mock_state->selected_item = 4;
@@ -537,12 +540,64 @@ static void menu_action_tests(void)
 		screen_state_eq(mock_state, MOCK_SCREEN_MENU, 1, MOCK_IGNORE);
 	}
 
-	/* vb2_ui_back_action */
+	VB2_DEBUG("...done.\n");
+}
+
+static void try_recovery_action_tests(void)
+{
+	VB2_DEBUG("Testing try recovery action...\n");
+
+	/* SUCCESS */
+	reset_common_data();
+	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS,
+		"SUCCESS");
+	TEST_EQ(mock_get_screen_info_called, 0, "  no change_screen");
+
+	/* NO_DISK_FOUND */
+	reset_common_data();
+	set_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"NO_DISK_FOUND");
+	screen_state_eq(mock_state, VB2_SCREEN_RECOVERY_SELECT,
+			MOCK_IGNORE, MOCK_IGNORE);
+
+	/* NO_DISK_FOUND -> INVALID_KERNEL -> SUCCESS */
+	reset_common_data();
+	set_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"NO_DISK_FOUND");
+	set_mock_vbtlk(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+		       VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"INVALID_KERNEL");
+	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS, "SUCCESS");
+	screen_state_eq(mock_state, VB2_SCREEN_RECOVERY_INVALID,
+			MOCK_IGNORE, MOCK_IGNORE);
+
+	/* INVALID_KERNEL */
 	reset_common_data();
-	mock_ui_context.key = VB_KEY_ESC;
-	TEST_EQ(vb2_ui_back_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"vb2_ui_back_action");
-	screen_state_eq(mock_state, VB2_SCREEN_BLANK, 0, MOCK_IGNORE);
+	set_mock_vbtlk(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+		       VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"INVALID_KERNEL");
+	screen_state_eq(mock_state, VB2_SCREEN_RECOVERY_INVALID,
+			MOCK_IGNORE, MOCK_IGNORE);
+
+	/* INVALID_KERNEL -> NO_DISK_FOUND -> SUCCESS */
+	reset_common_data();
+	set_mock_vbtlk(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+		       VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"INVALID_KERNEL");
+	set_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"NO_DISK_FOUND");
+	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS, "SUCCESS");
+	screen_state_eq(mock_state, VB2_SCREEN_RECOVERY_SELECT,
+			MOCK_IGNORE, MOCK_IGNORE);
 
 	VB2_DEBUG("...done.\n");
 }
@@ -569,13 +624,13 @@ static void ui_loop_tests(void)
 	/* Global action */
 	reset_common_data();
 	mock_calls_until_shutdown = -1;
-	TEST_EQ(ui_loop(ctx, VB2_SCREEN_BLANK, global_action_countdown),
+	TEST_EQ(ui_loop(ctx, VB2_SCREEN_BLANK, mock_action_countdown),
 		VB2_SUCCESS, "global action");
-	TEST_EQ(global_action_called, 10, "  global action called");
+	TEST_EQ(mock_action_called, 10, "  action called");
 
 	/* Global action can change screen */
 	reset_common_data();
-	TEST_EQ(ui_loop(ctx, VB2_SCREEN_BLANK, global_action_change_screen),
+	TEST_EQ(ui_loop(ctx, VB2_SCREEN_BLANK, mock_action_change_screen),
 		VB2_REQUEST_SHUTDOWN, "global action can change screen");
 	displayed_eq("pass", MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE,
 		     MOCK_IGNORE);
@@ -591,6 +646,7 @@ static void ui_loop_tests(void)
 	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_DOWN);  /* (blocked) */
 	add_mock_keypress(VB_KEY_UP);
+	add_mock_keypress(VB_KEY_UP);
 	add_mock_keypress(VB_KEY_ENTER);
 	TEST_EQ(ui_loop(ctx, MOCK_SCREEN_MENU, NULL),
 		VB2_REQUEST_SHUTDOWN, "KEY_UP, KEY_DOWN, and KEY_ENTER");
@@ -606,7 +662,9 @@ static void ui_loop_tests(void)
 		     MOCK_IGNORE);
 	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 3,
 		     MOCK_IGNORE);
-	displayed_eq("mock_screen_target_3", MOCK_SCREEN_TARGET3, MOCK_IGNORE,
+	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 2,
+		     MOCK_IGNORE);
+	displayed_eq("mock_screen_target_2", MOCK_SCREEN_TARGET2, MOCK_IGNORE,
 		     MOCK_IGNORE, MOCK_IGNORE);
 	displayed_no_extra();
 
@@ -620,6 +678,7 @@ static void ui_loop_tests(void)
 		add_mock_keypress(VB_BUTTON_VOL_DOWN_SHORT_PRESS);
 		add_mock_keypress(VB_BUTTON_VOL_DOWN_SHORT_PRESS);
 		add_mock_keypress(VB_BUTTON_VOL_UP_SHORT_PRESS);
+		add_mock_keypress(VB_BUTTON_VOL_UP_SHORT_PRESS);
 		add_mock_keypress(VB_BUTTON_POWER_SHORT_PRESS);
 		TEST_EQ(ui_loop(ctx, MOCK_SCREEN_MENU, NULL),
 			VB2_REQUEST_SHUTDOWN, "DETACHABLE");
@@ -635,7 +694,9 @@ static void ui_loop_tests(void)
 			     4, MOCK_IGNORE);
 		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
 			     3, MOCK_IGNORE);
-		displayed_eq("mock_screen_target_3", MOCK_SCREEN_TARGET3,
+		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
+			     2, MOCK_IGNORE);
+		displayed_eq("mock_screen_target_2", MOCK_SCREEN_TARGET2,
 			     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 		displayed_no_extra();
 	}
@@ -643,70 +704,18 @@ static void ui_loop_tests(void)
 	VB2_DEBUG("...done.\n");
 }
 
-static void try_recovery_action_tests(void)
+int main(void)
 {
-	VB2_DEBUG("Testing try recovery action...\n");
-
-	/* SUCCESS */
-	reset_common_data();
-	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS,
-		"SUCCESS");
-	TEST_EQ(mock_get_screen_info_last, -1, "  no change_screen");
-
-	/* NO_DISK_FOUND */
-	reset_common_data();
-	set_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"NO_DISK_FOUND");
-	TEST_EQ(mock_get_screen_info_last, VB2_SCREEN_RECOVERY_SELECT,
-		"  recovery select screen");
+	/* Input actions */
+	menu_up_action_tests();
+	menu_down_action_tests();
+	menu_select_action_tests();
 
-	/* NO_DISK_FOUND -> INVALID_KERNEL -> SUCCESS */
-	reset_common_data();
-	set_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"NO_DISK_FOUND");
-	set_mock_vbtlk(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       VB_DISK_FLAG_REMOVABLE);
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"INVALID_KERNEL");
-	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS, "SUCCESS");
-	TEST_EQ(mock_get_screen_info_last, VB2_SCREEN_RECOVERY_INVALID,
-		"  recovery invalid screen");
-
-	/* INVALID_KERNEL */
-	reset_common_data();
-	set_mock_vbtlk(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       VB_DISK_FLAG_REMOVABLE);
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"INVALID_KERNEL");
-	TEST_EQ(mock_get_screen_info_last, VB2_SCREEN_RECOVERY_INVALID,
-		"  recovery invalid screen");
-
-	/* INVALID_KERNEL -> NO_DISK_FOUND -> SUCCESS */
-	reset_common_data();
-	set_mock_vbtlk(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       VB_DISK_FLAG_REMOVABLE);
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"INVALID_KERNEL");
-	set_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"NO_DISK_FOUND");
-	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS, "SUCCESS");
-	TEST_EQ(mock_get_screen_info_last, VB2_SCREEN_RECOVERY_SELECT,
-		"  recovery select screen");
-
-	VB2_DEBUG("...done.\n");
-}
+	/* Global actions */
+	try_recovery_action_tests();
 
-int main(void)
-{
-	menu_action_tests();
+	/* Core UI loop */
 	ui_loop_tests();
-	try_recovery_action_tests();
 
 	return gTestSuccess ? 0 : 255;
 }
diff --git a/tests/vb2_ui_tests.c b/tests/vb2_ui_tests.c
index abb05f9c..13c037c1 100644
--- a/tests/vb2_ui_tests.c
+++ b/tests/vb2_ui_tests.c
@@ -165,18 +165,8 @@ static void reset_common_data(enum reset_type t)
 	sd->status |= VB2_SD_STATUS_SECDATA_KERNEL_INIT;
 
 	/* Mock ui_context based on real screens */
-	mock_ui_context = (struct vb2_ui_context){
-		.ctx = ctx,
-		.root_screen = vb2_get_screen_info(VB2_SCREEN_BLANK),
-		.state = (struct vb2_screen_state){
-			.screen = vb2_get_screen_info(VB2_SCREEN_BLANK),
-			.selected_item = 0,
-			.disabled_item_mask = 0,
-		},
-		.locale_id = 0,
-		.key = 0,
-
-	};
+	memset(&mock_ui_context, 0, sizeof(mock_ui_context));
+	mock_ui_context.ctx = ctx;
 	mock_state = &mock_ui_context.state;
 
 	/* For vb2ex_display_ui */
diff --git a/tests/vb2_ui_utility_tests.c b/tests/vb2_ui_utility_tests.c
index a3952617..0a7e4593 100644
--- a/tests/vb2_ui_utility_tests.c
+++ b/tests/vb2_ui_utility_tests.c
@@ -18,9 +18,10 @@
 #define MOCK_IGNORE 0xffffu
 
 /* Mock screen index for testing screen utility functions. */
-#define MOCK_NO_SCREEN 0xef0
-#define MOCK_SCREEN_BASE 0xeff
-#define MOCK_SCREEN_MENU 0xfff
+#define MOCK_NO_SCREEN 0xef00
+#define MOCK_SCREEN_BASE 0xef10
+#define MOCK_SCREEN_MENU 0xef11
+#define MOCK_SCREEN_ROOT 0xefff
 
 /* Mock data */
 static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
@@ -70,6 +71,12 @@ const struct vb2_screen_info mock_screen_menu = {
 	.num_items = ARRAY_SIZE(mock_screen_menu_items),
 	.items = mock_screen_menu_items,
 };
+const struct vb2_screen_info mock_screen_root = {
+	.id = MOCK_SCREEN_ROOT,
+	.name = "mock_screen_root",
+	.num_items = ARRAY_SIZE(mock_empty_menu),
+	.items = mock_empty_menu,
+};
 
 static void screen_state_eq(const struct vb2_screen_state *state,
 			    enum vb2_screen screen,
@@ -104,19 +111,8 @@ static void reset_common_data(void)
 	mock_shutdown_request = MOCK_IGNORE;
 
 	/* Mock ui_context based on mock screens */
-	mock_ui_context = (struct vb2_ui_context){
-		.ctx = ctx,
-		.root_screen = &mock_screen_blank,
-		.state = (struct vb2_screen_state){
-			.screen = &mock_screen_blank,
-			.selected_item = 0,
-			.disabled_item_mask = 0,
-		},
-		.locale_id = 0,
-		.key = 0,
-		.power_button = VB2_POWER_BUTTON_HELD_SINCE_BOOT,
-
-	};
+	memset(&mock_ui_context, 0, sizeof(mock_ui_context));
+	mock_ui_context.power_button = VB2_POWER_BUTTON_HELD_SINCE_BOOT;
 	mock_state = &mock_ui_context.state;
 }
 
@@ -143,6 +139,8 @@ const struct vb2_screen_info *vb2_get_screen_info(enum vb2_screen screen)
 		return &mock_screen_base;
 	case MOCK_SCREEN_MENU:
 		return &mock_screen_menu;
+	case MOCK_SCREEN_ROOT:
+		return &mock_screen_root;
 	default:
 		return NULL;
 	}
@@ -246,6 +244,22 @@ static void check_shutdown_request_tests(void)
 	VB2_DEBUG("...done.\n");
 }
 
+static void vb2_ui_back_action_tests(void)
+{
+	VB2_DEBUG("Testing vb2_ui_back_action...\n");
+
+	/* TODO: back to previous screen instead of root screen */
+	/* Back to root screen */
+	reset_common_data();
+	mock_ui_context.root_screen = &mock_screen_root;
+	mock_ui_context.key = VB_KEY_ESC;
+	TEST_EQ(vb2_ui_back_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"back to root screen");
+	screen_state_eq(mock_state, MOCK_SCREEN_ROOT, MOCK_IGNORE, MOCK_IGNORE);
+
+	VB2_DEBUG("...done.\n");
+}
+
 static void change_screen_tests(void)
 {
 	VB2_DEBUG("Testing change_screen...\n");
@@ -276,6 +290,7 @@ static void change_screen_tests(void)
 int main(void)
 {
 	check_shutdown_request_tests();
+	vb2_ui_back_action_tests();
 	change_screen_tests();
 
 	return gTestSuccess ? 0 : 255;