vboot/ui: Split UI tests

Rearrange UI tests into three files: vb2_ui_utility_tests for utility
functions and core UI functions, vb2_ui_action_tests for hooked actions,
and vb2_ui_tests for UI entries.

This CL does not add any extra tests.

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

Signed-off-by: Hsuan Ting Chen <roccochen@chromium.org>
Change-Id: I36a0a43aa3295b06cf32446dcc107652d64d2b8f
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2198268
Reviewed-by: Joel Kitching <kitching@chromium.org>

4 files changed, 751 insertions, 567 deletions

diff --git a/Makefile b/Makefile
index 10c5c4bf..b35dc913 100644
--- a/Makefile
+++ b/Makefile
@@ -724,6 +724,7 @@ TEST2X_NAMES = \
 	tests/vb2_sha_api_tests \
 	tests/vb2_sha_tests \
 	tests/vb2_ui_tests \
+	tests/vb2_ui_action_tests \
 	tests/vb2_ui_utility_tests \
 	tests/hmac_test
 
@@ -1249,6 +1250,7 @@ run2tests: install_for_test
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_sha_api_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_sha_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_ui_tests
+	${RUNTEST} ${BUILD_RUN}/tests/vb2_ui_action_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_ui_utility_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb20_api_kernel_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb20_kernel_tests
diff --git a/tests/vb2_ui_action_tests.c b/tests/vb2_ui_action_tests.c
new file mode 100644
index 00000000..dd87bc33
--- /dev/null
+++ b/tests/vb2_ui_action_tests.c
@@ -0,0 +1,734 @@
+/* Copyright 2020 The Chromium OS Authors. All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ *
+ * Tests for UI related actions.
+ */
+
+#include "2api.h"
+#include "2common.h"
+#include "2misc.h"
+#include "2nvstorage.h"
+#include "2ui.h"
+#include "2ui_private.h"
+#include "test_common.h"
+#include "vboot_kernel.h"
+
+/* Fixed value for ignoring some checks. */
+#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
+
+/* Mock data */
+struct display_call {
+	const struct vb2_screen_info *screen;
+	uint32_t locale_id;
+	uint32_t selected_item;
+	uint32_t disabled_item_mask;
+};
+
+static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
+	__attribute__((aligned(VB2_WORKBUF_ALIGN)));
+static struct vb2_context *ctx;
+static struct vb2_gbb_header gbb;
+
+static int mock_calls_until_shutdown;
+
+static struct vb2_ui_context mock_ui_context;
+static struct vb2_screen_state *mock_state;
+
+static struct display_call mock_displayed[64];
+static int mock_displayed_count;
+static int mock_displayed_i;
+
+static uint32_t mock_key[64];
+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 vb2_error_t mock_vbtlk_retval;
+static uint32_t mock_vbtlk_expected_flag;
+
+/* Mock screens */
+const struct vb2_menu_item mock_empty_menu[] = {};
+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 = {
+	.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[] = {
+	{
+		.text = "option 0",
+		.target = MOCK_SCREEN_TARGET0,
+	},
+	{
+		.text = "option 1",
+		.target = MOCK_SCREEN_TARGET1,
+	},
+	{
+		.text = "option 2",
+		.target = MOCK_SCREEN_TARGET2,
+	},
+	{
+		.text = "option 3",
+		.target = MOCK_SCREEN_TARGET3,
+	},
+	{
+		.text = "option 4 (no target)",
+	},
+};
+const struct vb2_screen_info mock_screen_menu = {
+	.id = MOCK_SCREEN_MENU,
+	.name = "mock_screen_menu: screen with 5 options",
+	.num_items = ARRAY_SIZE(mock_screen_menu_items),
+	.items = mock_screen_menu_items,
+};
+const struct vb2_screen_info mock_screen_target0 = {
+	.id = MOCK_SCREEN_TARGET0,
+	.name = "mock_screen_target0",
+	.num_items = ARRAY_SIZE(mock_empty_menu),
+	.items = mock_empty_menu,
+};
+const struct vb2_screen_info mock_screen_target1 = {
+	.id = MOCK_SCREEN_TARGET1,
+	.name = "mock_screen_target1",
+	.num_items = ARRAY_SIZE(mock_empty_menu),
+	.items = mock_empty_menu,
+};
+const struct vb2_screen_info mock_screen_target2 = {
+	.id = MOCK_SCREEN_TARGET2,
+	.name = "mock_screen_target2",
+	.num_items = ARRAY_SIZE(mock_empty_menu),
+	.items = mock_empty_menu,
+};
+const struct vb2_screen_info mock_screen_target3 = {
+	.id = MOCK_SCREEN_TARGET3,
+	.name = "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)
+{
+	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,
+			    uint32_t selected_item,
+			    uint32_t disabled_item_mask)
+{
+	if (screen != MOCK_IGNORE) {
+		if (state->screen == NULL)
+			TEST_TRUE(0, "  state.screen does not exist");
+		else
+			TEST_EQ(state->screen->id, screen, "  state.screen");
+	}
+	if (selected_item != MOCK_IGNORE)
+		TEST_EQ(state->selected_item,
+			selected_item, "  state.selected_item");
+	if (disabled_item_mask != MOCK_IGNORE)
+		TEST_EQ(state->disabled_item_mask,
+			disabled_item_mask, "  state.disabled_item_mask");
+}
+
+static void add_mock_key(uint32_t press, int trusted)
+{
+	if (mock_key_total >= ARRAY_SIZE(mock_key) ||
+	    mock_key_total >= ARRAY_SIZE(mock_key_trusted)) {
+		TEST_TRUE(0, "  mock_key ran out of entries!");
+		return;
+	}
+
+	mock_key[mock_key_total] = press;
+	mock_key_trusted[mock_key_total] = trusted;
+	mock_key_total++;
+}
+
+static void add_mock_keypress(uint32_t press)
+{
+	add_mock_key(press, 0);
+}
+
+
+static void set_mock_vbtlk(vb2_error_t retval, uint32_t get_info_flags)
+{
+	mock_vbtlk_retval = retval;
+	mock_vbtlk_expected_flag = get_info_flags;
+}
+
+static void displayed_eq(const char *text,
+			 enum vb2_screen screen,
+			 uint32_t locale_id,
+			 uint32_t selected_item,
+			 uint32_t disabled_item_mask)
+{
+	char text_buf[256];
+
+	if (mock_displayed_i >= mock_displayed_count) {
+		sprintf(text_buf, "  missing screen %s", text);
+		TEST_TRUE(0, text_buf);
+		return;
+	}
+
+	if (screen != MOCK_IGNORE) {
+		sprintf(text_buf, "  screen of %s", text);
+		TEST_EQ(mock_displayed[mock_displayed_i].screen->id, screen,
+			text_buf);
+	}
+	if (locale_id != MOCK_IGNORE) {
+		sprintf(text_buf, "  locale_id of %s", text);
+		TEST_EQ(mock_displayed[mock_displayed_i].locale_id, locale_id,
+			text_buf);
+	}
+	if (selected_item != MOCK_IGNORE) {
+		sprintf(text_buf, "  selected_item of %s", text);
+		TEST_EQ(mock_displayed[mock_displayed_i].selected_item,
+			selected_item, text_buf);
+	}
+	if (disabled_item_mask != MOCK_IGNORE) {
+		sprintf(text_buf, "  disabled_item_mask of %s", text);
+		TEST_EQ(mock_displayed[mock_displayed_i].disabled_item_mask,
+			disabled_item_mask, text_buf);
+	}
+	mock_displayed_i++;
+}
+
+static void displayed_no_extra(void)
+{
+	if (mock_displayed_i == 0)
+		TEST_EQ(mock_displayed_count, 0, "  no screen");
+	else
+		TEST_EQ(mock_displayed_count, mock_displayed_i,
+			"  no extra screens");
+}
+
+/* Reset mock data (for use before each test) */
+static void reset_common_data(void)
+{
+	TEST_SUCC(vb2api_init(workbuf, sizeof(workbuf), &ctx),
+		  "vb2api_init failed");
+
+	memset(&gbb, 0, sizeof(gbb));
+
+	vb2_nv_init(ctx);
+
+	/* For shutdown_required */
+	mock_calls_until_shutdown = 10;
+
+	/* For try_recovery_action */
+	invalid_disk_last = -1;
+
+	/* 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,
+
+	};
+	mock_state = &mock_ui_context.state;
+
+	/* For vb2ex_display_ui */
+	memset(mock_displayed, 0, sizeof(mock_displayed));
+	mock_displayed_count = 0;
+	mock_displayed_i = 0;
+
+	/* 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;
+
+	/* For global actions */
+	global_action_called = 0;
+
+	/* For vb2_get_screen_info */
+	mock_get_screen_info_last = -1;
+
+	/* For VbTryLoadKernel */
+	mock_vbtlk_retval = VB2_ERROR_MOCK;
+	mock_vbtlk_expected_flag = MOCK_IGNORE;
+}
+
+/* Mock functions */
+struct vb2_gbb_header *vb2_get_gbb(struct vb2_context *c)
+{
+	return &gbb;
+}
+
+uint32_t VbExIsShutdownRequested(void)
+{
+	if (mock_calls_until_shutdown < 0)  /* Never request shutdown */
+		return 0;
+	if (mock_calls_until_shutdown == 0)
+		return 1;
+	mock_calls_until_shutdown--;
+
+	return 0;
+}
+
+const struct vb2_screen_info *vb2_get_screen_info(enum vb2_screen screen)
+{
+	mock_get_screen_info_last = screen;
+
+	switch ((int)screen) {
+	case VB2_SCREEN_BLANK:
+		return &mock_screen_blank;
+	case MOCK_SCREEN_BASE:
+		return &mock_screen_base;
+	case MOCK_SCREEN_MENU:
+		return &mock_screen_menu;
+	case MOCK_SCREEN_TARGET0:
+		return &mock_screen_target0;
+	case MOCK_SCREEN_TARGET1:
+		return &mock_screen_target1;
+	case MOCK_SCREEN_TARGET2:
+		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;
+	}
+}
+
+vb2_error_t vb2ex_display_ui(enum vb2_screen screen,
+			     uint32_t locale_id,
+			     uint32_t selected_item,
+			     uint32_t disabled_item_mask)
+{
+	VB2_DEBUG("displayed %d: screen = %#x, locale_id = %u, "
+		  "selected_item = %u, disabled_item_mask = %#x\n",
+		  mock_displayed_count, screen, locale_id, selected_item,
+		  disabled_item_mask);
+
+	if (mock_displayed_count >= ARRAY_SIZE(mock_displayed)) {
+		TEST_TRUE(0, "  mock vb2ex_display_ui ran out of entries!");
+		return VB2_ERROR_MOCK;
+	}
+
+	mock_displayed[mock_displayed_count] = (struct display_call){
+		.screen = vb2_get_screen_info(screen),
+		.locale_id = locale_id,
+		.selected_item = selected_item,
+		.disabled_item_mask = disabled_item_mask,
+	};
+	mock_displayed_count++;
+
+	return VB2_SUCCESS;
+}
+
+uint32_t VbExKeyboardRead(void)
+{
+	return VbExKeyboardReadWithFlags(NULL);
+}
+
+uint32_t VbExKeyboardReadWithFlags(uint32_t *key_flags)
+{
+	if (mock_key_count < mock_key_total) {
+		if (key_flags != NULL) {
+			if (mock_key_trusted[mock_key_count])
+				*key_flags = VB_KEY_FLAG_TRUSTED_KEYBOARD;
+			else
+				*key_flags = 0;
+		}
+		return mock_key[mock_key_count++];
+	}
+
+	return 0;
+}
+
+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)
+{
+	int i, target_id;
+	char test_name[256];
+
+	VB2_DEBUG("Testing menu actions...\n");
+
+	/* Valid menu_up_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");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 1, MOCK_IGNORE);
+
+	/* Valid menu_up_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");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 0, MOCK_IGNORE);
+
+	/* Invalid menu_up_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)");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 0, MOCK_IGNORE);
+
+	/* Invalid menu_up_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)");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
+
+	/* Ignore volume-up when not DETACHABLE */
+	if (!DETACHABLE) {
+		reset_common_data();
+		mock_state->screen = &mock_screen_menu;
+		mock_state->selected_item = 2;
+		mock_ui_context.key = VB_BUTTON_VOL_UP_SHORT_PRESS;
+		TEST_EQ(menu_up_action(&mock_ui_context),
+			VB2_REQUEST_UI_CONTINUE,
+			"ignore volume-up when not DETACHABLE");
+		screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
+	}
+
+	/* Valid menu_down_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");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 3, MOCK_IGNORE);
+
+	/* Valid menu_down_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");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 4, MOCK_IGNORE);
+
+	/* Invalid menu_down_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)");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 4, MOCK_IGNORE);
+
+	/* Invalid menu_down_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)");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
+
+	/* Ignore volume-down when not DETACHABLE */
+	if (!DETACHABLE) {
+		reset_common_data();
+		mock_state->screen = &mock_screen_menu;
+		mock_state->selected_item = 2;
+		mock_ui_context.key = VB_BUTTON_VOL_DOWN_SHORT_PRESS;
+		TEST_EQ(menu_down_action(&mock_ui_context),
+			VB2_REQUEST_UI_CONTINUE,
+			"ignore volume-down when not DETACHABLE");
+		screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
+	}
+
+	/* menu_select_action with no item screen */
+	reset_common_data();
+	mock_state->screen = &mock_screen_base;
+	mock_ui_context.key = VB_KEY_ENTER;
+	TEST_EQ(menu_select_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"menu_select_action with no item screen");
+	screen_state_eq(mock_state, MOCK_SCREEN_BASE, 0, MOCK_IGNORE);
+
+	/* Try to select target 0..3 */
+	for (i = 0; i <= 3; i++) {
+		sprintf(test_name, "select target %d", i);
+		target_id = MOCK_SCREEN_TARGET0 + i;
+		reset_common_data();
+		mock_state->screen = &mock_screen_menu;
+		mock_state->selected_item = i;
+		mock_ui_context.key = VB_KEY_ENTER;
+		TEST_EQ(menu_select_action(&mock_ui_context),
+			VB2_REQUEST_UI_CONTINUE, test_name);
+		screen_state_eq(mock_state, target_id, 0, MOCK_IGNORE);
+	}
+
+	/* Try to select no target item */
+	reset_common_data();
+	mock_state->screen = &mock_screen_menu;
+	mock_state->selected_item = 4;
+	mock_ui_context.key = VB_KEY_ENTER;
+	TEST_EQ(menu_select_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"select no target");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 4, MOCK_IGNORE);
+
+	/* Ignore power button short press when not DETACHABLE */
+	if (!DETACHABLE) {
+		reset_common_data();
+		mock_state->screen = &mock_screen_menu;
+		mock_state->selected_item = 1;
+		mock_ui_context.key = VB_BUTTON_POWER_SHORT_PRESS;
+		TEST_EQ(menu_select_action(&mock_ui_context),
+			VB2_REQUEST_UI_CONTINUE,
+			"ignore power button short press when not DETACHABLE");
+		screen_state_eq(mock_state, MOCK_SCREEN_MENU, 1, MOCK_IGNORE);
+	}
+
+	/* vb2_ui_back_action */
+	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);
+
+	VB2_DEBUG("...done.\n");
+}
+
+static void ui_loop_tests(void)
+{
+	VB2_DEBUG("Testing ui_loop...\n");
+
+	/* Die if no root screen */
+	reset_common_data();
+	TEST_ABORT(ui_loop(ctx, MOCK_NO_SCREEN, NULL),
+		   "die if no root screen");
+	displayed_no_extra();
+
+	/* Shutdown if requested */
+	reset_common_data();
+	TEST_EQ(ui_loop(ctx, MOCK_SCREEN_BASE, NULL),
+		VB2_REQUEST_SHUTDOWN, "shutdown if requested");
+	TEST_EQ(mock_calls_until_shutdown, 0, "  used up shutdown request");
+	displayed_eq("mock_screen_base", MOCK_SCREEN_BASE, MOCK_IGNORE,
+		     MOCK_IGNORE, MOCK_IGNORE);
+	displayed_no_extra();
+
+	/* Global action */
+	reset_common_data();
+	mock_calls_until_shutdown = -1;
+	TEST_EQ(ui_loop(ctx, VB2_SCREEN_BLANK, global_action_countdown),
+		VB2_SUCCESS, "global action");
+	TEST_EQ(global_action_called, 10, "  global action called");
+
+	/* Global action can change screen */
+	reset_common_data();
+	TEST_EQ(ui_loop(ctx, VB2_SCREEN_BLANK, global_action_change_screen),
+		VB2_REQUEST_SHUTDOWN, "global action can change screen");
+	displayed_eq("pass", MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE,
+		     MOCK_IGNORE);
+	displayed_eq("change to mock_screen_base", MOCK_IGNORE, MOCK_IGNORE,
+		     MOCK_IGNORE, MOCK_IGNORE);
+
+	/* KEY_UP, KEY_DOWN, and KEY_ENTER navigation */
+	reset_common_data();
+	add_mock_keypress(VB_KEY_UP);  /* (blocked) */
+	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_DOWN);  /* (blocked) */
+	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");
+	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 0,
+		     MOCK_IGNORE);
+	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 1,
+		     MOCK_IGNORE);
+	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 2,
+		     MOCK_IGNORE);
+	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 3,
+		     MOCK_IGNORE);
+	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 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, MOCK_IGNORE,
+		     MOCK_IGNORE, MOCK_IGNORE);
+	displayed_no_extra();
+
+	/* For DETACHABLE */
+	if (DETACHABLE) {
+		reset_common_data();
+		add_mock_keypress(VB_BUTTON_VOL_UP_SHORT_PRESS);
+		add_mock_keypress(VB_BUTTON_VOL_DOWN_SHORT_PRESS);
+		add_mock_keypress(VB_BUTTON_VOL_DOWN_SHORT_PRESS);
+		add_mock_keypress(VB_BUTTON_VOL_DOWN_SHORT_PRESS);
+		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_POWER_SHORT_PRESS);
+		TEST_EQ(ui_loop(ctx, MOCK_SCREEN_MENU, NULL),
+			VB2_REQUEST_SHUTDOWN, "DETACHABLE");
+		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
+			     0, MOCK_IGNORE);
+		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
+			     1, MOCK_IGNORE);
+		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
+			     2, MOCK_IGNORE);
+		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
+			     3, MOCK_IGNORE);
+		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
+			     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,
+			     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+		displayed_no_extra();
+	}
+
+	VB2_DEBUG("...done.\n");
+}
+
+static void try_recovery_action_tests(void)
+{
+	VB2_DEBUG("Testing try recovery action...\n");
+
+	/* Success on the first try */
+	reset_common_data();
+	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS,
+		"success on the first try");
+	TEST_EQ(mock_get_screen_info_last, -1, "  no change_screen");
+
+	/* No disk found on the first try */
+	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 on the first try");
+	TEST_EQ(mock_get_screen_info_last, VB2_SCREEN_RECOVERY_SELECT,
+		"  recovery select screen");
+
+	/* Invalid disk on the first try */
+	reset_common_data();
+	set_mock_vbtlk(VB2_ERROR_MOCK, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"invalid on the first try");
+	TEST_EQ(mock_get_screen_info_last, VB2_SCREEN_RECOVERY_INVALID,
+		"  recovery invalid screen");
+
+	/* Success, last == 0 */
+	reset_common_data();
+	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
+	invalid_disk_last = 0;
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS,
+		"success, last == 0");
+	TEST_EQ(mock_get_screen_info_last, -1, "  no change_screen");
+
+	/* No disk found, last == 0 */
+	reset_common_data();
+	set_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	invalid_disk_last = 0;
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"no disk found, last == 0");
+	TEST_EQ(mock_get_screen_info_last, -1, "  no change_screen");
+
+	/* Invalid disk, last == 0 */
+	reset_common_data();
+	set_mock_vbtlk(VB2_ERROR_MOCK, VB_DISK_FLAG_REMOVABLE);
+	invalid_disk_last = 0;
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"invalid, last == 0");
+	TEST_EQ(mock_get_screen_info_last, VB2_SCREEN_RECOVERY_INVALID,
+		"  recovery invalid screen");
+
+	/* Success, last == 1 */
+	reset_common_data();
+	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
+	invalid_disk_last = 1;
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS,
+		"success, last == 1");
+	TEST_EQ(mock_get_screen_info_last, -1, "  no change_screen");
+
+	/* No disk found, last == 1 */
+	reset_common_data();
+	set_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	invalid_disk_last = 1;
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"no disk found, last == 1");
+	TEST_EQ(mock_get_screen_info_last, VB2_SCREEN_RECOVERY_SELECT,
+		"  recovery select screen");
+
+	/* Invalid disk, last == 1 */
+	reset_common_data();
+	set_mock_vbtlk(VB2_ERROR_MOCK, VB_DISK_FLAG_REMOVABLE);
+	invalid_disk_last = 1;
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"invalid, last == 1");
+	TEST_EQ(mock_get_screen_info_last, -1, "  no change_screen");
+
+	VB2_DEBUG("...done.\n");
+}
+
+int main(void)
+{
+	menu_action_tests();
+	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 c7605d74..2f652880 100644
--- a/tests/vb2_ui_tests.c
+++ b/tests/vb2_ui_tests.c
@@ -51,7 +51,7 @@ static int mock_dev_boot_legacy_allowed;
 static int mock_dev_boot_usb_allowed;
 
 static int mock_vbexlegacy_called;
-static enum VbAltFwIndex_t mock_altfw_num;
+static enum VbAltFwIndex_t mock_altfw_num_last;
 
 static vb2_error_t mock_vbtlk_retval[32];
 static uint32_t mock_vbtlk_expected_flag[32];
@@ -147,8 +147,10 @@ static void reset_common_data(void)
 
 	sd = vb2_get_sd(ctx);
 
-	/* For global actions */
+	/* For try_recovery_action */
 	invalid_disk_last = -1;
+
+	/* 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),
@@ -187,7 +189,7 @@ static void reset_common_data(void)
 
 	/* For VbExLegacy */
 	mock_vbexlegacy_called = 0;
-	mock_altfw_num = -100;
+	mock_altfw_num_last = -100;
 
 	/* For VbTryLoadKernel */
 	memset(mock_vbtlk_retval, 0, sizeof(mock_vbtlk_retval));
@@ -282,7 +284,7 @@ int vb2_dev_boot_usb_allowed(struct vb2_context *c)
 vb2_error_t VbExLegacy(enum VbAltFwIndex_t altfw_num)
 {
 	mock_vbexlegacy_called++;
-	mock_altfw_num = altfw_num;
+	mock_altfw_num_last = altfw_num;
 
 	return VB2_SUCCESS;
 }
@@ -305,89 +307,6 @@ vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t get_info_flags)
 }
 
 /* Tests */
-static void try_recovery_action_tests(void)
-{
-	VB2_DEBUG("Testing try recovery action...\n");
-
-	/* Success on the first try */
-	reset_common_data();
-	add_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS,
-		"success on the first try");
-	TEST_EQ(mock_state->screen->id, VB2_SCREEN_BLANK,
-		"  screen remains the same");
-
-	/* No disk found on the first try */
-	reset_common_data();
-	add_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 on the first try");
-	TEST_EQ(mock_state->screen->id, VB2_SCREEN_RECOVERY_SELECT,
-		"  recovery select screen");
-
-	/* Invalid disk on the first try */
-	reset_common_data();
-	add_mock_vbtlk(VB2_ERROR_MOCK, VB_DISK_FLAG_REMOVABLE);
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"invalid on the first try");
-	TEST_EQ(mock_state->screen->id, VB2_SCREEN_RECOVERY_INVALID,
-		"  recovery invalid screen");
-
-	/* Success, last == 0 */
-	reset_common_data();
-	add_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
-	invalid_disk_last = 0;
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS,
-		"success, last == 0");
-	TEST_EQ(mock_state->screen->id, VB2_SCREEN_BLANK,
-		"  screen remains the same");
-
-	/* No disk found, last == 0 */
-	reset_common_data();
-	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
-	invalid_disk_last = 0;
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"no disk found, last == 0");
-	TEST_EQ(mock_state->screen->id, VB2_SCREEN_BLANK, "  screen no change");
-
-	/* Invalid disk, last == 0 */
-	reset_common_data();
-	add_mock_vbtlk(VB2_ERROR_MOCK, VB_DISK_FLAG_REMOVABLE);
-	invalid_disk_last = 0;
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"invalid, last == 0");
-	TEST_EQ(mock_state->screen->id, VB2_SCREEN_RECOVERY_INVALID,
-		"  recovery invalid screen");
-
-	/* Success, last == 1 */
-	reset_common_data();
-	add_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
-	invalid_disk_last = 1;
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS,
-		"success, last == 1");
-	TEST_EQ(mock_state->screen->id, VB2_SCREEN_BLANK,
-		"  screen remains the same");
-
-	/* No disk found, last == 1 */
-	reset_common_data();
-	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
-	invalid_disk_last = 1;
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"no disk found, last == 1");
-	TEST_EQ(mock_state->screen->id, VB2_SCREEN_RECOVERY_SELECT,
-		"  recovery select screen");
-
-	/* Invalid disk, last == 1 */
-	reset_common_data();
-	add_mock_vbtlk(VB2_ERROR_MOCK, VB_DISK_FLAG_REMOVABLE);
-	invalid_disk_last = 1;
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"invalid, last == 1");
-	TEST_EQ(mock_state->screen->id, VB2_SCREEN_BLANK, "  screen no change");
-
-	VB2_DEBUG("...done.\n");
-}
-
 static void developer_tests(void)
 {
 	VB2_DEBUG("Testing developer mode...\n");
@@ -407,7 +326,7 @@ static void developer_tests(void)
 	mock_dev_boot_legacy_allowed = 1;
 	TEST_EQ(vb2_developer_menu(ctx), VB2_SUCCESS, "proceed to legacy");
 	TEST_EQ(mock_vbexlegacy_called, 1, "  try legacy");
-	TEST_EQ(mock_altfw_num, 0, "  check altfw_num");
+	TEST_EQ(mock_altfw_num_last, 0, "  check altfw_num");
 	displayed_no_extra();
 	TEST_EQ(mock_vbtlk_count, mock_vbtlk_total, "  used up mock_vbtlk");
 
@@ -575,7 +494,6 @@ static void manual_recovery_tests(void)
 
 int main(void)
 {
-	try_recovery_action_tests();
 	developer_tests();
 	broken_recovery_tests();
 	manual_recovery_tests();
diff --git a/tests/vb2_ui_utility_tests.c b/tests/vb2_ui_utility_tests.c
index 08b537f5..5350e1c6 100644
--- a/tests/vb2_ui_utility_tests.c
+++ b/tests/vb2_ui_utility_tests.c
@@ -2,7 +2,7 @@
  * Use of this source code is governed by a BSD-style license that can be
  * found in the LICENSE file.
  *
- * Tests for UI functions without real UI flow.
+ * Tests for UI utility functions.
  */
 
 #include "2api.h"
@@ -13,7 +13,6 @@
 #include "2ui_private.h"
 #include "test_common.h"
 #include "vboot_api.h"
-#include "vboot_kernel.h"
 
 /* Fixed value for ignoring some checks. */
 #define MOCK_IGNORE 0xffffu
@@ -22,20 +21,8 @@
 #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
 
 /* Mock data */
-struct display_call {
-	const struct vb2_screen_info *screen;
-	uint32_t locale_id;
-	uint32_t selected_item;
-	uint32_t disabled_item_mask;
-};
-
 static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
 	__attribute__((aligned(VB2_WORKBUF_ALIGN)));
 static struct vb2_context *ctx;
@@ -46,16 +33,7 @@ static int mock_shutdown_request;
 static struct vb2_ui_context mock_ui_context;
 static struct vb2_screen_state *mock_state;
 
-static struct display_call mock_displayed[64];
-static int mock_displayed_count;
-static int mock_displayed_i;
-
-static uint32_t mock_key[64];
-static int mock_key_trusted[64];
-static int mock_key_count;
-static int mock_key_total;
-
-/* Mocks for testing screen utility functions. */
+/* Mock screens */
 const struct vb2_menu_item mock_empty_menu[] = {};
 struct vb2_screen_info mock_screen_blank = {
 	.id = VB2_SCREEN_BLANK,
@@ -63,94 +41,35 @@ struct vb2_screen_info mock_screen_blank = {
 	.num_items = ARRAY_SIZE(mock_empty_menu),
 	.items = mock_empty_menu,
 };
-struct vb2_screen_info mock_screen_base =
-{
+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[] =
-{
+struct vb2_menu_item mock_screen_menu_items[] = {
 	{
 		.text = "option 0",
-		.target = MOCK_SCREEN_TARGET0,
 	},
 	{
 		.text = "option 1",
-		.target = MOCK_SCREEN_TARGET1,
 	},
 	{
 		.text = "option 2",
-		.target = MOCK_SCREEN_TARGET2,
 	},
 	{
 		.text = "option 3",
-		.target = MOCK_SCREEN_TARGET3,
 	},
 	{
-		.text = "option 4 (no target)",
+		.text = "option 4",
 	},
 };
-const struct vb2_screen_info mock_screen_menu =
-{
+const struct vb2_screen_info mock_screen_menu = {
 	.id = MOCK_SCREEN_MENU,
 	.name = "mock_screen_menu: screen with 5 options",
 	.num_items = ARRAY_SIZE(mock_screen_menu_items),
 	.items = mock_screen_menu_items,
 };
-const struct vb2_screen_info mock_screen_target0 =
-{
-	.id = MOCK_SCREEN_TARGET0,
-	.name = "mock_screen_target0",
-	.num_items = ARRAY_SIZE(mock_empty_menu),
-	.items = mock_empty_menu,
-};
-const struct vb2_screen_info mock_screen_target1 =
-{
-	.id = MOCK_SCREEN_TARGET1,
-	.name = "mock_screen_target1",
-	.num_items = ARRAY_SIZE(mock_empty_menu),
-	.items = mock_empty_menu,
-};
-const struct vb2_screen_info mock_screen_target2 =
-{
-	.id = MOCK_SCREEN_TARGET2,
-	.name = "mock_screen_target2",
-	.num_items = ARRAY_SIZE(mock_empty_menu),
-	.items = mock_empty_menu,
-};
-const struct vb2_screen_info mock_screen_target3 =
-{
-	.id = MOCK_SCREEN_TARGET3,
-	.name = "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,
-};
-
-/* Actions for tests */
-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)
-{
-	change_screen(ui, MOCK_SCREEN_BASE);
-	if (++global_action_called >= 10)
-		return VB2_SUCCESS;
-	return VB2_REQUEST_UI_CONTINUE;
-}
 
 static void screen_state_eq(const struct vb2_screen_state *state,
 			    enum vb2_screen screen,
@@ -171,70 +90,6 @@ static void screen_state_eq(const struct vb2_screen_state *state,
 			disabled_item_mask, "  state.disabled_item_mask");
 }
 
-static void add_mock_key(uint32_t press, int trusted)
-{
-	if (mock_key_total >= ARRAY_SIZE(mock_key) ||
-	    mock_key_total >= ARRAY_SIZE(mock_key_trusted)) {
-		TEST_TRUE(0, "  mock_key ran out of entries!");
-		return;
-	}
-
-	mock_key[mock_key_total] = press;
-	mock_key_trusted[mock_key_total] = trusted;
-	mock_key_total++;
-}
-
-static void add_mock_keypress(uint32_t press)
-{
-	add_mock_key(press, 0);
-}
-
-static void displayed_eq(const char *text,
-			 enum vb2_screen screen,
-			 uint32_t locale_id,
-			 uint32_t selected_item,
-			 uint32_t disabled_item_mask)
-{
-	char text_buf[256];
-
-	if (mock_displayed_i >= mock_displayed_count) {
-		sprintf(text_buf, "  missing screen %s", text);
-		TEST_TRUE(0, text_buf);
-		return;
-	}
-
-	if (screen != MOCK_IGNORE) {
-		sprintf(text_buf, "  screen of %s", text);
-		TEST_EQ(mock_displayed[mock_displayed_i].screen->id, screen,
-			text_buf);
-	}
-	if (locale_id != MOCK_IGNORE) {
-		sprintf(text_buf, "  locale_id of %s", text);
-		TEST_EQ(mock_displayed[mock_displayed_i].locale_id, locale_id,
-			text_buf);
-	}
-	if (selected_item != MOCK_IGNORE) {
-		sprintf(text_buf, "  selected_item of %s", text);
-		TEST_EQ(mock_displayed[mock_displayed_i].selected_item,
-			selected_item, text_buf);
-	}
-	if (disabled_item_mask != MOCK_IGNORE) {
-		sprintf(text_buf, "  disabled_item_mask of %s", text);
-		TEST_EQ(mock_displayed[mock_displayed_i].disabled_item_mask,
-			disabled_item_mask, text_buf);
-	}
-	mock_displayed_i++;
-}
-
-static void displayed_no_extra(void)
-{
-	if (mock_displayed_i == 0)
-		TEST_EQ(mock_displayed_count, 0, "  no screen");
-	else
-		TEST_EQ(mock_displayed_count, mock_displayed_i,
-			"  no extra screens");
-}
-
 /* Reset mock data (for use before each test) */
 static void reset_common_data(void)
 {
@@ -249,7 +104,7 @@ static void reset_common_data(void)
 	power_button = POWER_BUTTON_HELD_SINCE_BOOT;
 	mock_shutdown_request = MOCK_IGNORE;
 
-	/* For menu actions */
+	/* Mock ui_context based on mock screens */
 	mock_ui_context = (struct vb2_ui_context){
 		.ctx = ctx,
 		.root_screen = &mock_screen_blank,
@@ -263,20 +118,6 @@ static void reset_common_data(void)
 
 	};
 	mock_state = &mock_ui_context.state;
-
-	/* For vb2ex_display_ui */
-	memset(mock_displayed, 0, sizeof(mock_displayed));
-	mock_displayed_count = 0;
-	mock_displayed_i = 0;
-
-	/* 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;
-
-	/* For global actions */
-	global_action_called = 0;
 }
 
 /* Mock functions */
@@ -302,67 +143,11 @@ 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_TARGET0:
-		return &mock_screen_target0;
-	case MOCK_SCREEN_TARGET1:
-		return &mock_screen_target1;
-	case MOCK_SCREEN_TARGET2:
-		return &mock_screen_target2;
-	case MOCK_SCREEN_TARGET3:
-		return &mock_screen_target3;
-	case MOCK_SCREEN_TARGET4:
-		return &mock_screen_target4;
 	default:
 		return NULL;
 	}
 }
 
-vb2_error_t vb2ex_display_ui(enum vb2_screen screen,
-			     uint32_t locale_id,
-			     uint32_t selected_item,
-			     uint32_t disabled_item_mask)
-{
-	VB2_DEBUG("displayed %d: screen = %#x, locale_id = %u, "
-		  "selected_item = %u, disabled_item_mask = %#x\n",
-		  mock_displayed_count, screen, locale_id, selected_item,
-		  disabled_item_mask);
-
-	if (mock_displayed_count >= ARRAY_SIZE(mock_displayed)) {
-		TEST_TRUE(0, "  mock vb2ex_display_ui ran out of entries!");
-		return VB2_ERROR_MOCK;
-	}
-
-	mock_displayed[mock_displayed_count] = (struct display_call){
-		.screen = vb2_get_screen_info(screen),
-		.locale_id = locale_id,
-		.selected_item = selected_item,
-		.disabled_item_mask = disabled_item_mask,
-	};
-	mock_displayed_count++;
-
-	return VB2_SUCCESS;
-}
-
-uint32_t VbExKeyboardRead(void)
-{
-	return VbExKeyboardReadWithFlags(NULL);
-}
-
-uint32_t VbExKeyboardReadWithFlags(uint32_t *key_flags)
-{
-	if (mock_key_count < mock_key_total) {
-		if (key_flags != NULL) {
-			if (mock_key_trusted[mock_key_count])
-				*key_flags = VB_KEY_FLAG_TRUSTED_KEYBOARD;
-			else
-				*key_flags = 0;
-		}
-		return mock_key[mock_key_count++];
-	}
-
-	return 0;
-}
-
 /* Tests */
 static void shutdown_required_tests(void)
 {
@@ -446,165 +231,6 @@ static void shutdown_required_tests(void)
 	VB2_DEBUG("...done.\n");
 }
 
-static void menu_action_tests(void)
-{
-	int i, target_id;
-	char test_name[256];
-
-	VB2_DEBUG("Testing menu actions...\n");
-
-	/* Valid menu_up_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");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 1, MOCK_IGNORE);
-
-	/* Valid menu_up_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");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 0, MOCK_IGNORE);
-
-	/* Invalid menu_up_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)");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 0, MOCK_IGNORE);
-
-	/* Invalid menu_up_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)");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
-
-	/* Ignore volume-up when not DETACHABLE */
-	if (!DETACHABLE) {
-		reset_common_data();
-		mock_state->screen = &mock_screen_menu;
-		mock_state->selected_item = 2;
-		mock_ui_context.key = VB_BUTTON_VOL_UP_SHORT_PRESS;
-		TEST_EQ(menu_up_action(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE,
-			"ignore volume-up when not DETACHABLE");
-		screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
-	}
-
-	/* Valid menu_down_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");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 3, MOCK_IGNORE);
-
-	/* Valid menu_down_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");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 4, MOCK_IGNORE);
-
-	/* Invalid menu_down_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)");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 4, MOCK_IGNORE);
-
-	/* Invalid menu_down_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)");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
-
-	/* Ignore volume-down when not DETACHABLE */
-	if (!DETACHABLE) {
-		reset_common_data();
-		mock_state->screen = &mock_screen_menu;
-		mock_state->selected_item = 2;
-		mock_ui_context.key = VB_BUTTON_VOL_DOWN_SHORT_PRESS;
-		TEST_EQ(menu_down_action(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE,
-			"ignore volume-down when not DETACHABLE");
-		screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
-	}
-
-	/* menu_select_action with no item screen */
-	reset_common_data();
-	mock_state->screen = &mock_screen_base;
-	mock_ui_context.key = VB_KEY_ENTER;
-	TEST_EQ(menu_select_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"menu_select_action with no item screen");
-	screen_state_eq(mock_state, MOCK_SCREEN_BASE, 0, MOCK_IGNORE);
-
-	/* Try to select target 0..3 */
-	for (i = 0; i <= 3; i++) {
-		sprintf(test_name, "select target %d", i);
-		target_id = MOCK_SCREEN_TARGET0 + i;
-		reset_common_data();
-		mock_state->screen = &mock_screen_menu;
-		mock_state->selected_item = i;
-		mock_ui_context.key = VB_KEY_ENTER;
-		TEST_EQ(menu_select_action(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE, test_name);
-		screen_state_eq(mock_state, target_id, 0, MOCK_IGNORE);
-	}
-
-	/* Try to select no target item */
-	reset_common_data();
-	mock_state->screen = &mock_screen_menu;
-	mock_state->selected_item = 4;
-	mock_ui_context.key = VB_KEY_ENTER;
-	TEST_EQ(menu_select_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"select no target");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 4, MOCK_IGNORE);
-
-	/* Ignore power button short press when not DETACHABLE */
-	if (!DETACHABLE) {
-		reset_common_data();
-		mock_state->screen = &mock_screen_menu;
-		mock_state->selected_item = 1;
-		mock_ui_context.key = VB_BUTTON_POWER_SHORT_PRESS;
-		TEST_EQ(menu_select_action(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE,
-			"ignore power button short press when not DETACHABLE");
-		screen_state_eq(mock_state, MOCK_SCREEN_MENU, 1, MOCK_IGNORE);
-	}
-
-	/* vb2_ui_back_action */
-	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);
-
-	VB2_DEBUG("...done.\n");
-}
-
 static void change_screen_tests(void)
 {
 	VB2_DEBUG("Testing change_screen...\n");
@@ -627,101 +253,7 @@ static void change_screen_tests(void)
 		"change to screen which does not exist");
 	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, MOCK_IGNORE);
 
-	VB2_DEBUG("...done.\n");
-}
-
-static void ui_loop_tests(void)
-{
-	VB2_DEBUG("Testing ui_loop...\n");
-
-	/* Die if no root screen */
-	reset_common_data();
-	TEST_ABORT(ui_loop(ctx, MOCK_NO_SCREEN, NULL),
-		   "die if no root screen");
-	displayed_no_extra();
-
-	/* Shutdown if requested */
-	reset_common_data();
-	mock_shutdown_request = VB_SHUTDOWN_REQUEST_OTHER;
-	TEST_EQ(ui_loop(ctx, MOCK_SCREEN_BASE, NULL),
-		VB2_REQUEST_SHUTDOWN, "shutdown if requested");
-	displayed_eq("mock_screen_base", MOCK_SCREEN_BASE, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE);
-	displayed_no_extra();
-
-	/* Global action */
-	reset_common_data();
-	TEST_EQ(ui_loop(ctx, VB2_SCREEN_BLANK, global_action_countdown),
-		VB2_SUCCESS, "global action");
-	TEST_EQ(global_action_called, 10, "  global action called");
-
-	/* Global action can change screen */
-	reset_common_data();
-	TEST_EQ(ui_loop(ctx, VB2_SCREEN_BLANK, global_action_change_screen),
-		VB2_SUCCESS, "global action can change screen");
-	TEST_EQ(global_action_called, 10, "  global action called");
-	displayed_eq("pass", MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE,
-		     MOCK_IGNORE);
-	displayed_eq("change to mock_screen_base", MOCK_IGNORE, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE);
-
-	/* KEY_UP, KEY_DOWN, and KEY_ENTER navigation */
-	reset_common_data();
-	add_mock_keypress(VB_KEY_UP);  /* (blocked) */
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_DOWN);  /* (blocked) */
-	add_mock_keypress(VB_KEY_UP);
-	add_mock_keypress(VB_KEY_ENTER);
-	TEST_EQ(ui_loop(ctx, MOCK_SCREEN_MENU, global_action_countdown),
-		VB2_SUCCESS, "KEY_UP, KEY_DOWN, and KEY_ENTER");
-	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 0,
-		     MOCK_IGNORE);
-	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 1,
-		     MOCK_IGNORE);
-	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 2,
-		     MOCK_IGNORE);
-	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 3,
-		     MOCK_IGNORE);
-	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 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, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE);
-	displayed_no_extra();
-
-	/* For DETACHABLE */
-	if (DETACHABLE) {
-		reset_common_data();
-		add_mock_keypress(VB_BUTTON_VOL_UP_SHORT_PRESS);
-		add_mock_keypress(VB_BUTTON_VOL_DOWN_SHORT_PRESS);
-		add_mock_keypress(VB_BUTTON_VOL_DOWN_SHORT_PRESS);
-		add_mock_keypress(VB_BUTTON_VOL_DOWN_SHORT_PRESS);
-		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_POWER_SHORT_PRESS);
-		TEST_EQ(ui_loop(ctx, MOCK_SCREEN_MENU, global_action_countdown),
-			VB2_SUCCESS, "DETACHABLE");
-		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
-			     0, MOCK_IGNORE);
-		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
-			     1, MOCK_IGNORE);
-		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
-			     2, MOCK_IGNORE);
-		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
-			     3, MOCK_IGNORE);
-		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
-			     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,
-			     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-		displayed_no_extra();
-	}
+	/* TODO: Change to screen with init */
 
 	VB2_DEBUG("...done.\n");
 }
@@ -729,9 +261,7 @@ static void ui_loop_tests(void)
 int main(void)
 {
 	shutdown_required_tests();
-	menu_action_tests();
 	change_screen_tests();
-	ui_loop_tests();
 
 	return gTestSuccess ? 0 : 255;
 }