vboot/ui: Remove diagnostics menu and related tests

This CL is a part of centralizing ui codes. The removed screens and unit
tests will be added in depthcharge.

Remove diagnostics menu, diagnostics storage and memory checking screens.

Remove unit tests for above screens and menu.

BUG=b:172339016
TEST=export CC=x86_64-pc-linux-gnu-clang DEBUG=1 DETACHABLE=0; \
     make -j test_setup && make -j runtests
TEST=export CC=x86_64-pc-linux-gnu-clang DEBUG=1 DETACHABLE=1; \
     make -j test_setup && make -j runtests
BRANCH=none

Cq-Depend: chromium:3138687
Signed-off-by: edisonhello <edisonhello@google.com>
Change-Id: I5eb1bee6a5f1aedb77298acbddab20156c1c086e
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/3139537
Reviewed-by: Yu-Ping Wu <yupingso@chromium.org>

9 files changed, 19 insertions, 1033 deletions

diff --git a/Makefile b/Makefile
index fd13991e..be95aeba 100644
--- a/Makefile
+++ b/Makefile
@@ -755,7 +755,6 @@ TEST2X_NAMES = \
 	tests/vb2_secdata_kernel_tests \
 	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
@@ -1298,7 +1297,6 @@ run2tests: install_for_test
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_secdata_kernel_tests
 	${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
diff --git a/firmware/2lib/2stub.c b/firmware/2lib/2stub.c
index 181fa496..47fb2df7 100644
--- a/firmware/2lib/2stub.c
+++ b/firmware/2lib/2stub.c
@@ -274,3 +274,9 @@ vb2_error_t vb2ex_developer_ui(struct vb2_context *ctx)
 {
 	return VB2_SUCCESS;
 }
+
+__attribute__((weak))
+vb2_error_t vb2ex_diagnostic_ui(struct vb2_context *ctx)
+{
+	return VB2_SUCCESS;
+}
diff --git a/firmware/2lib/2ui.c b/firmware/2lib/2ui.c
index ab0692bf..64680220 100644
--- a/firmware/2lib/2ui.c
+++ b/firmware/2lib/2ui.c
@@ -412,11 +412,3 @@ vb2_error_t vb2_ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 		return VB2_SUCCESS;
 	return rv;
 }
-
-/*****************************************************************************/
-/* Diagnostics */
-
-vb2_error_t vb2_diagnostic_menu(struct vb2_context *ctx)
-{
-	return vb2_ui_loop(ctx, VB2_SCREEN_DIAGNOSTICS, NULL);
-}
diff --git a/firmware/2lib/2ui_screens.c b/firmware/2lib/2ui_screens.c
index 5e08b57d..3a8e0b60 100644
--- a/firmware/2lib/2ui_screens.c
+++ b/firmware/2lib/2ui_screens.c
@@ -125,25 +125,6 @@ static vb2_error_t log_page_reset_to_top(struct vb2_ui_context *ui)
 	return log_page_update(ui, NULL);
 }
 
-static vb2_error_t log_page_show_back_or_cancel(struct vb2_ui_context *ui,
-						int is_show_cancel)
-{
-	int back_item = ui->state->screen->back_item;
-	int cancel_item = ui->state->screen->cancel_item;
-	VB2_CLR_BIT(ui->state->hidden_item_mask, back_item);
-	VB2_CLR_BIT(ui->state->hidden_item_mask, cancel_item);
-	if (is_show_cancel) {
-		VB2_SET_BIT(ui->state->hidden_item_mask, back_item);
-		if (ui->state->selected_item == back_item)
-			ui->state->selected_item = cancel_item;
-	} else {
-		VB2_SET_BIT(ui->state->hidden_item_mask, cancel_item);
-		if (ui->state->selected_item == cancel_item)
-			ui->state->selected_item = back_item;
-	}
-	return VB2_SUCCESS;
-}
-
 static vb2_error_t log_page_prev_action(struct vb2_ui_context *ui)
 {
 	/* Validity check. */
@@ -514,331 +495,6 @@ static const struct vb2_screen_info recovery_to_dev_screen = {
 	.menu = MENU_ITEMS(recovery_to_dev_items),
 };
 
-/******************************************************************************/
-/* VB2_SCREEN_DIAGNOSTICS */
-
-#define DIAGNOSTICS_ITEM_STORAGE_HEALTH 1
-#define DIAGNOSTICS_ITEM_STORAGE_TEST_SHORT 2
-#define DIAGNOSTICS_ITEM_STORAGE_TEST_EXTENDED 3
-
-static vb2_error_t diagnostics_init(struct vb2_ui_context *ui)
-{
-	const char *unused_log_string;
-	vb2_error_t rv = vb2ex_diag_get_storage_test_log(&unused_log_string);
-	if (rv == VB2_ERROR_EX_UNIMPLEMENTED) {
-		VB2_SET_BIT(ui->state->disabled_item_mask,
-			    DIAGNOSTICS_ITEM_STORAGE_TEST_SHORT);
-		VB2_SET_BIT(ui->state->disabled_item_mask,
-			    DIAGNOSTICS_ITEM_STORAGE_TEST_EXTENDED);
-	}
-	ui->state->selected_item = DIAGNOSTICS_ITEM_STORAGE_HEALTH;
-	return VB2_SUCCESS;
-}
-
-static const struct vb2_menu_item diagnostics_items[] = {
-	LANGUAGE_SELECT_ITEM,
-	[DIAGNOSTICS_ITEM_STORAGE_HEALTH] = {
-		.text = "Storage health info",
-		.target = VB2_SCREEN_DIAGNOSTICS_STORAGE_HEALTH,
-	},
-	[DIAGNOSTICS_ITEM_STORAGE_TEST_SHORT] = {
-		.text = "Storage self-test (short)",
-		.target = VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST_SHORT,
-	},
-	[DIAGNOSTICS_ITEM_STORAGE_TEST_EXTENDED] = {
-		.text = "Storage self-test (Extended)",
-		.target = VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST_EXTENDED,
-	},
-	{
-		.text = "Memory check (quick)",
-		.target = VB2_SCREEN_DIAGNOSTICS_MEMORY_QUICK,
-	},
-	{
-		.text = "Memory check (full)",
-		.target = VB2_SCREEN_DIAGNOSTICS_MEMORY_FULL,
-	},
-	POWER_OFF_ITEM,
-};
-
-static const struct vb2_screen_info diagnostics_screen = {
-	.id = VB2_SCREEN_DIAGNOSTICS,
-	.name = "Diagnostic tools",
-	.init = diagnostics_init,
-	.menu = MENU_ITEMS(diagnostics_items),
-};
-
-/******************************************************************************/
-/* VB2_SCREEN_DIAGNOSTICS_STORAGE_HEALTH */
-
-#define DIAGNOSTICS_STORAGE_HEALTH_ITEM_PAGE_UP 0
-#define DIAGNOSTICS_STORAGE_HEALTH_ITEM_PAGE_DOWN 1
-#define DIAGNOSTICS_STORAGE_HEALTH_ITEM_BACK 2
-
-static vb2_error_t diagnostics_storage_health_init_impl(
-	struct vb2_ui_context *ui)
-{
-	const char *log_string;
-	VB2_TRY(vb2ex_diag_get_storage_health(&log_string));
-	VB2_TRY(log_page_update(ui, log_string));
-	return log_page_reset_to_top(ui);
-}
-
-static vb2_error_t diagnostics_storage_health_init(struct vb2_ui_context *ui)
-{
-	if (vb2_is_error(diagnostics_storage_health_init_impl(ui)))
-		return set_ui_error_and_go_back(ui, VB2_UI_ERROR_DIAGNOSTICS);
-	return VB2_SUCCESS;
-}
-
-static const struct vb2_menu_item diagnostics_storage_health_items[] = {
-	[DIAGNOSTICS_STORAGE_HEALTH_ITEM_PAGE_UP] = PAGE_UP_ITEM,
-	[DIAGNOSTICS_STORAGE_HEALTH_ITEM_PAGE_DOWN] = PAGE_DOWN_ITEM,
-	[DIAGNOSTICS_STORAGE_HEALTH_ITEM_BACK] = BACK_ITEM,
-	POWER_OFF_ITEM,
-};
-
-static const struct vb2_screen_info diagnostics_storage_health_screen = {
-	.id = VB2_SCREEN_DIAGNOSTICS_STORAGE_HEALTH,
-	.name = "Storage health info",
-	.init = diagnostics_storage_health_init,
-	.menu = MENU_ITEMS(diagnostics_storage_health_items),
-	.page_up_item = DIAGNOSTICS_STORAGE_HEALTH_ITEM_PAGE_UP,
-	.page_down_item = DIAGNOSTICS_STORAGE_HEALTH_ITEM_PAGE_DOWN,
-	.back_item = DIAGNOSTICS_STORAGE_HEALTH_ITEM_BACK,
-};
-
-/******************************************************************************/
-/* VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST */
-
-#define DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_UP 0
-#define DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_DOWN 1
-#define DIAGNOSTICS_STORAGE_TEST_ITEM_BACK 2
-#define DIAGNOSTICS_STORAGE_TEST_ITEM_CANCEL 3
-
-static vb2_error_t diagnostics_storage_test_update_impl(
-	struct vb2_ui_context *ui)
-{
-	const char *log_string;
-	int is_test_running = 0;
-
-	/* Early return if the test is done. */
-	if (ui->state->test_finished)
-		return VB2_SUCCESS;
-
-	vb2_error_t rv = vb2ex_diag_get_storage_test_log(&log_string);
-	switch (rv) {
-	case VB2_ERROR_EX_DIAG_TEST_RUNNING:
-		is_test_running = 1;
-		break;
-	case VB2_SUCCESS:
-		ui->state->test_finished = 1;
-		break;
-	default:
-		VB2_DEBUG("vb2ex_diag_get_storage_test_log returned %#x\n", rv);
-		return rv;
-	}
-	VB2_TRY(log_page_show_back_or_cancel(ui, is_test_running));
-	return log_page_update(ui, log_string);
-}
-
-static vb2_error_t diagnostics_storage_test_update(struct vb2_ui_context *ui)
-{
-	if (vb2_is_error(diagnostics_storage_test_update_impl(ui)))
-		return set_ui_error_and_go_back(ui, VB2_UI_ERROR_DIAGNOSTICS);
-	return VB2_SUCCESS;
-}
-
-static vb2_error_t diagnostics_storage_test_control(
-	struct vb2_ui_context *ui, enum vb2_diag_storage_test op)
-{
-	if (vb2_is_error(vb2ex_diag_storage_test_control(op)))
-		return set_ui_error_and_go_back(ui, VB2_UI_ERROR_DIAGNOSTICS);
-	return VB2_SUCCESS;
-}
-
-static vb2_error_t diagnostics_storage_test_init(struct vb2_ui_context *ui)
-{
-	VB2_TRY(diagnostics_storage_test_update(ui));
-	if (vb2_is_error(log_page_reset_to_top(ui)))
-		return set_ui_error_and_go_back(ui, VB2_UI_ERROR_DIAGNOSTICS);
-	return VB2_SUCCESS;
-}
-
-static vb2_error_t diagnostics_storage_test_short_init(
-	struct vb2_ui_context *ui)
-{
-	VB2_TRY(diagnostics_storage_test_control(ui,
-						 VB2_DIAG_STORAGE_TEST_STOP));
-	VB2_TRY(diagnostics_storage_test_control(ui,
-						 VB2_DIAG_STORAGE_TEST_SHORT));
-	return diagnostics_storage_test_init(ui);
-}
-
-static vb2_error_t diagnostics_storage_test_extended_init(
-	struct vb2_ui_context *ui)
-{
-	VB2_TRY(diagnostics_storage_test_control(ui,
-						 VB2_DIAG_STORAGE_TEST_STOP));
-	VB2_TRY(diagnostics_storage_test_control(
-		ui, VB2_DIAG_STORAGE_TEST_EXTENDED));
-	return diagnostics_storage_test_init(ui);
-}
-
-static vb2_error_t diagnostics_storage_test_cancel(struct vb2_ui_context *ui)
-{
-	VB2_TRY(diagnostics_storage_test_control(ui,
-						 VB2_DIAG_STORAGE_TEST_STOP));
-	return vb2_ui_screen_back(ui);
-}
-
-static const struct vb2_menu_item diagnostics_storage_test_items[] = {
-	[DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_UP] = PAGE_UP_ITEM,
-	[DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_DOWN] = PAGE_DOWN_ITEM,
-	[DIAGNOSTICS_STORAGE_TEST_ITEM_BACK] = BACK_ITEM,
-	[DIAGNOSTICS_STORAGE_TEST_ITEM_CANCEL] = {
-		.text = "Cancel",
-		.action = diagnostics_storage_test_cancel,
-	},
-	POWER_OFF_ITEM,
-};
-
-static const struct vb2_screen_info diagnostics_storage_test_short_screen = {
-	.id = VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST_SHORT,
-	.name = "Storage self-test (short)",
-	.init = diagnostics_storage_test_short_init,
-	.action = diagnostics_storage_test_update,
-	.menu = MENU_ITEMS(diagnostics_storage_test_items),
-	.page_up_item = DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_UP,
-	.page_down_item = DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_DOWN,
-	.back_item = DIAGNOSTICS_STORAGE_TEST_ITEM_BACK,
-	.cancel_item = DIAGNOSTICS_STORAGE_TEST_ITEM_CANCEL,
-};
-
-static const struct vb2_screen_info diagnostics_storage_test_extended_screen = {
-	.id = VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST_EXTENDED,
-	.name = "Storage self-test (extended)",
-	.init = diagnostics_storage_test_extended_init,
-	.action = diagnostics_storage_test_update,
-	.menu = MENU_ITEMS(diagnostics_storage_test_items),
-	.page_up_item = DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_UP,
-	.page_down_item = DIAGNOSTICS_STORAGE_TEST_ITEM_PAGE_DOWN,
-	.back_item = DIAGNOSTICS_STORAGE_TEST_ITEM_BACK,
-	.cancel_item = DIAGNOSTICS_STORAGE_TEST_ITEM_CANCEL,
-};
-
-/******************************************************************************/
-/* VB2_SCREEN_DIAGNOSTICS_MEMORY_QUICK
-   VB2_SCREEN_DIAGNOSTICS_MEMORY_FULL */
-
-#define DIAGNOSTICS_MEMORY_ITEM_PAGE_UP 0
-#define DIAGNOSTICS_MEMORY_ITEM_PAGE_DOWN 1
-#define DIAGNOSTICS_MEMORY_ITEM_BACK 2
-#define DIAGNOSTICS_MEMORY_ITEM_CANCEL 3
-
-typedef vb2_error_t (*memory_test_op_t)(int reset, const char **out);
-static vb2_error_t diagnostics_memory_update_screen_impl(
-	struct vb2_ui_context *ui, memory_test_op_t op, int reset)
-{
-	const char *log_string = NULL;
-	vb2_error_t rv;
-	int is_test_running = 0;
-
-	/* Early return if the memory test is done. */
-	if (ui->state->test_finished)
-		return VB2_SUCCESS;
-
-	rv = op(reset, &log_string);
-	switch (rv) {
-	/* The test is still running but the output buffer was unchanged. */
-	case VB2_ERROR_EX_DIAG_TEST_RUNNING:
-		return VB2_SUCCESS;
-	case VB2_ERROR_EX_DIAG_TEST_UPDATED:
-		is_test_running = 1;
-		break;
-	case VB2_SUCCESS:
-		ui->state->test_finished = 1;
-		break;
-	default:
-		VB2_DEBUG("memory_test_op returned %#x\n", rv);
-		return rv;
-	}
-	VB2_TRY(log_page_show_back_or_cancel(ui, is_test_running));
-	return log_page_update(ui, log_string);
-}
-
-static vb2_error_t diagnostics_memory_update_screen(struct vb2_ui_context *ui,
-						    memory_test_op_t op,
-						    int reset)
-{
-	if (vb2_is_error(diagnostics_memory_update_screen_impl(ui, op, reset)))
-		return set_ui_error_and_go_back(ui, VB2_UI_ERROR_DIAGNOSTICS);
-	return VB2_SUCCESS;
-}
-
-static vb2_error_t diagnostics_memory_init_quick(struct vb2_ui_context *ui)
-{
-	VB2_TRY(diagnostics_memory_update_screen(
-		ui, &vb2ex_diag_memory_quick_test, 1));
-	if (vb2_is_error(log_page_reset_to_top(ui)))
-		return set_ui_error_and_go_back(ui, VB2_UI_ERROR_DIAGNOSTICS);
-	return VB2_SUCCESS;
-}
-
-static vb2_error_t diagnostics_memory_init_full(struct vb2_ui_context *ui)
-{
-	VB2_TRY(diagnostics_memory_update_screen(
-		ui, &vb2ex_diag_memory_full_test, 1));
-	if (vb2_is_error(log_page_reset_to_top(ui)))
-		return set_ui_error_and_go_back(ui, VB2_UI_ERROR_DIAGNOSTICS);
-	return VB2_SUCCESS;
-}
-
-static vb2_error_t diagnostics_memory_update_quick(struct vb2_ui_context *ui)
-{
-	return diagnostics_memory_update_screen(
-		ui, &vb2ex_diag_memory_quick_test, 0);
-}
-
-static vb2_error_t diagnostics_memory_update_full(struct vb2_ui_context *ui)
-{
-	return diagnostics_memory_update_screen(
-		ui, &vb2ex_diag_memory_full_test, 0);
-}
-
-static const struct vb2_menu_item diagnostics_memory_items[] = {
-	[DIAGNOSTICS_MEMORY_ITEM_PAGE_UP] = PAGE_UP_ITEM,
-	[DIAGNOSTICS_MEMORY_ITEM_PAGE_DOWN] = PAGE_DOWN_ITEM,
-	[DIAGNOSTICS_MEMORY_ITEM_BACK] = BACK_ITEM,
-	[DIAGNOSTICS_MEMORY_ITEM_CANCEL] = {
-		.text = "Cancel",
-		.action = vb2_ui_screen_back,
-	},
-	POWER_OFF_ITEM,
-};
-
-static const struct vb2_screen_info diagnostics_memory_quick_screen = {
-	.id = VB2_SCREEN_DIAGNOSTICS_MEMORY_QUICK,
-	.name = "Memory check (quick)",
-	.init = diagnostics_memory_init_quick,
-	.action = diagnostics_memory_update_quick,
-	.menu = MENU_ITEMS(diagnostics_memory_items),
-	.page_up_item = DIAGNOSTICS_MEMORY_ITEM_PAGE_UP,
-	.page_down_item = DIAGNOSTICS_MEMORY_ITEM_PAGE_DOWN,
-	.back_item = DIAGNOSTICS_MEMORY_ITEM_BACK,
-	.cancel_item = DIAGNOSTICS_MEMORY_ITEM_CANCEL,
-};
-
-static const struct vb2_screen_info diagnostics_memory_full_screen = {
-	.id = VB2_SCREEN_DIAGNOSTICS_MEMORY_FULL,
-	.name = "Memory check (full)",
-	.init = diagnostics_memory_init_full,
-	.action = diagnostics_memory_update_full,
-	.menu = MENU_ITEMS(diagnostics_memory_items),
-	.page_up_item = DIAGNOSTICS_MEMORY_ITEM_PAGE_UP,
-	.page_down_item = DIAGNOSTICS_MEMORY_ITEM_PAGE_DOWN,
-	.back_item = DIAGNOSTICS_MEMORY_ITEM_BACK,
-	.cancel_item = DIAGNOSTICS_MEMORY_ITEM_CANCEL,
-};
 
 /******************************************************************************/
 /*
@@ -854,12 +510,6 @@ static const struct vb2_screen_info *screens[] = {
 	&debug_info_screen,
 	&firmware_log_screen,
 	&recovery_to_dev_screen,
-	&diagnostics_screen,
-	&diagnostics_storage_health_screen,
-	&diagnostics_storage_test_short_screen,
-	&diagnostics_storage_test_extended_screen,
-	&diagnostics_memory_quick_screen,
-	&diagnostics_memory_full_screen,
 };
 
 const struct vb2_screen_info *vb2_get_screen_info(enum vb2_screen id)
diff --git a/firmware/2lib/include/2api.h b/firmware/2lib/include/2api.h
index a1cda6a1..521cae8f 100644
--- a/firmware/2lib/include/2api.h
+++ b/firmware/2lib/include/2api.h
@@ -1430,6 +1430,17 @@ vb2_error_t vb2ex_manual_recovery_ui(struct vb2_context *ctx);
  */
 vb2_error_t vb2ex_developer_ui(struct vb2_context *ctx);
 
+/**
+ * UI for a diagnostic tools boot.
+ *
+ * Enter the diagnostic tools menu, which provides debug information and
+ * diagnostic tests of various hardware components.
+ *
+ * @param ctx		Vboot context
+ * @return VB2_SUCCESS, or non-zero error code.
+ */
+vb2_error_t vb2ex_diagnostic_ui(struct vb2_context *ctx);
+
 /* Helpers for bitmask operations */
 #define VB2_SET_BIT(mask, index) ((mask) |= ((uint32_t)1 << (index)))
 #define VB2_CLR_BIT(mask, index) ((mask) &= ~((uint32_t)1 << (index)))
diff --git a/firmware/2lib/include/2ui.h b/firmware/2lib/include/2ui.h
index 0fb18513..c01c1f97 100644
--- a/firmware/2lib/include/2ui.h
+++ b/firmware/2lib/include/2ui.h
@@ -228,18 +228,4 @@ vb2_error_t vb2_ui_screen_back(struct vb2_ui_context *ui);
  */
 vb2_error_t vb2_ui_screen_change(struct vb2_ui_context *ui, enum vb2_screen id);
 
-/*****************************************************************************/
-/* UI loops */
-
-/**
- * UI for a diagnostic tools boot.
- *
- * Enter the diagnostic tools menu, which provides debug information and
- * diagnostic tests of various hardware components.
- *
- * @param ctx		Vboot context
- * @return VB2_SUCCESS, or non-zero error code.
- */
-vb2_error_t vb2_diagnostic_menu(struct vb2_context *ctx);
-
 #endif  /* VBOOT_REFERENCE_2UI_H_ */
diff --git a/firmware/lib/vboot_api_kernel.c b/firmware/lib/vboot_api_kernel.c
index 10728b95..b79b26fd 100644
--- a/firmware/lib/vboot_api_kernel.c
+++ b/firmware/lib/vboot_api_kernel.c
@@ -13,7 +13,6 @@
 #include "2rsa.h"
 #include "2secdata.h"
 #include "2sysincludes.h"
-#include "2ui.h"
 #include "load_kernel_fw.h"
 #include "vboot_api.h"
 #include "vboot_struct.h"
@@ -209,7 +208,7 @@ vb2_error_t VbSelectAndLoadKernel(struct vb2_context *ctx,
 		vb2ex_commit_data(ctx);
 
 		/* Diagnostic boot.  This has UI. */
-		VB2_TRY(vb2_diagnostic_menu(ctx));
+		VB2_TRY(vb2ex_diagnostic_ui(ctx));
 		/*
 		 * The diagnostic menu should either boot a rom, or
 		 * return either of reboot or shutdown.
diff --git a/tests/vb2_ui_tests.c b/tests/vb2_ui_tests.c
deleted file mode 100644
index b56d6315..00000000
--- a/tests/vb2_ui_tests.c
+++ /dev/null
@@ -1,655 +0,0 @@
-/* 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 developer and recovery mode UIs.
- */
-
-#include "2api.h"
-#include "2common.h"
-#include "2misc.h"
-#include "2nvstorage.h"
-#include "2struct.h"
-#include "2ui.h"
-#include "2ui_private.h"
-#include "test_common.h"
-#include "vboot_api.h"
-
-/* Fixed value for ignoring some checks */
-#define MOCK_IGNORE 0xffffu
-
-/* Fuzzy matches for check_time() */
-#define FUZZ_MS 200
-
-/* Mock data */
-/* TODO(b/156448738): Add tests for timer_disabled and error_code */
-struct display_call {
-	const struct vb2_screen_info *screen;
-	uint32_t locale_id;
-	uint32_t selected_item;
-	uint32_t disabled_item_mask;
-	uint32_t hidden_item_mask;
-	int timer_disabled;
-	uint32_t current_page;
-	enum vb2_ui_error error_code;
-} __attribute__((packed));
-
-struct beep_call {
-	uint32_t msec;
-	uint32_t frequency;
-	uint32_t time_expected;
-};
-
-static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
-	__attribute__((aligned(VB2_WORKBUF_ALIGN)));
-static struct vb2_context *ctx;
-static struct vb2_shared_data *sd;
-static struct vb2_gbb_header gbb;
-
-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_locale_count;
-
-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_total;
-
-static uint32_t mock_time_ms;
-static const uint32_t mock_time_start_ms = 31ULL * VB2_MSEC_PER_SEC;
-
-static struct beep_call mock_beep[8];
-static int mock_beep_count;
-static int mock_beep_total;
-
-static enum vb2_dev_default_boot_target mock_default_boot;
-
-static int mock_run_altfw_called;
-static uint32_t mock_altfw_last;
-static uint32_t mock_altfw_count;
-
-static vb2_error_t mock_vbtlk_retval[32];
-static uint32_t mock_vbtlk_expected_flag[32];
-static int mock_vbtlk_total;
-
-static int mock_allow_recovery;
-
-/* mock_pp_* = mock data for physical presence button */
-static int mock_pp_pressed[64];
-static int mock_pp_pressed_total;
-
-static int mock_enable_dev_mode;
-
-#define MOCK_PREPARE_LOG_SIZE 32
-
-static int mock_snapshot_count;
-static char mock_prepare_log[64][MOCK_PREPARE_LOG_SIZE];
-static int mock_prepare_log_count;
-static uint32_t mock_log_page_count;
-
-static vb2_error_t mock_diag_storage_test_rv;
-
-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 add_mock_vbtlk(vb2_error_t retval, uint32_t disk_flags)
-{
-	if (mock_vbtlk_total >= ARRAY_SIZE(mock_vbtlk_retval) ||
-	    mock_vbtlk_total >= ARRAY_SIZE(mock_vbtlk_expected_flag)) {
-		TEST_TRUE(0, "  mock_vbtlk ran out of entries!");
-		return;
-	}
-
-	mock_vbtlk_retval[mock_vbtlk_total] = retval;
-	mock_vbtlk_expected_flag[mock_vbtlk_total] = disk_flags;
-	mock_vbtlk_total++;
-}
-
-static void add_mock_pp_pressed(int pressed)
-{
-	if (mock_pp_pressed_total >= ARRAY_SIZE(mock_pp_pressed)) {
-		TEST_TRUE(0, "  mock_pp ran out of entries!");
-		return;
-	}
-
-	mock_pp_pressed[mock_pp_pressed_total++] = pressed;
-}
-
-static void displayed_eq(const char *text,
-			 enum vb2_screen screen,
-			 uint32_t locale_id,
-			 uint32_t selected_item,
-			 uint32_t disabled_item_mask,
-			 uint32_t hidden_item_mask,
-			 uint32_t current_page,
-			 int line)
-{
-	char text_info[32], text_buf[128];
-
-	sprintf(text_info, "(line #%d, displayed #%d)", line, mock_displayed_i);
-
-	if (mock_displayed_i >= mock_displayed_count) {
-		sprintf(text_buf, "  %s missing screen %s",
-			text_info, text);
-		TEST_TRUE(0, text_buf);
-		return;
-	}
-
-	if (screen != MOCK_IGNORE) {
-		sprintf(text_buf, "  %s screen of %s", text_info, text);
-		TEST_EQ(mock_displayed[mock_displayed_i].screen->id, screen,
-			text_buf);
-	}
-	if (locale_id != MOCK_IGNORE) {
-		sprintf(text_buf, "  %s locale_id of %s", text_info, text);
-		TEST_EQ(mock_displayed[mock_displayed_i].locale_id, locale_id,
-			text_buf);
-	}
-	if (selected_item != MOCK_IGNORE) {
-		sprintf(text_buf, "  %s selected_item of %s",
-			text_info, text);
-		TEST_EQ(mock_displayed[mock_displayed_i].selected_item,
-			selected_item, text_buf);
-	}
-	if (disabled_item_mask != MOCK_IGNORE) {
-		sprintf(text_buf, "  %s disabled_item_mask of %s",
-			text_info, text);
-		TEST_EQ(mock_displayed[mock_displayed_i].disabled_item_mask,
-			disabled_item_mask, text_buf);
-	}
-	if (hidden_item_mask != MOCK_IGNORE) {
-		sprintf(text_buf, "  %s hidden_item_mask of %s",
-			text_info, text);
-		TEST_EQ(mock_displayed[mock_displayed_i].hidden_item_mask,
-			hidden_item_mask, text_buf);
-	}
-	if (current_page != MOCK_IGNORE) {
-		sprintf(text_buf, "  %s current_page of %s",
-			text_info, text);
-		TEST_EQ(mock_displayed[mock_displayed_i].current_page,
-			current_page, text_buf);
-	}
-	mock_displayed_i++;
-}
-
-static void displayed_no_extra(int line)
-{
-	char text_info[32], text_buf[128];
-
-	sprintf(text_info, "(line #%d)", line);
-
-	if (mock_displayed_i == 0)
-		sprintf(text_buf, "  %s no screen", text_info);
-	else
-		sprintf(text_buf, "  %s no extra screens", text_info);
-	TEST_EQ(mock_displayed_count, mock_displayed_i, text_buf);
-}
-
-#define DISPLAYED_EQ(...) displayed_eq(__VA_ARGS__, __LINE__)
-
-#define DISPLAYED_PASS() \
-	displayed_eq("", MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, \
-		     MOCK_IGNORE, MOCK_IGNORE, __LINE__)
-
-#define DISPLAYED_NO_EXTRA() displayed_no_extra(__LINE__)
-
-/* Check if the result time falls in range [expected, expected + FUZZ_MS) */
-static void check_time(uint32_t result, uint32_t expected, const char *desc)
-{
-	TEST_TRUE(result >= expected, desc);
-	TEST_TRUE(result - expected < FUZZ_MS, "  within FUZZ_MS");
-}
-
-/* Type of test to reset for */
-enum reset_type {
-	FOR_DEVELOPER,
-	FOR_BROKEN_RECOVERY,
-	FOR_MANUAL_RECOVERY,
-	FOR_DIAGNOSTICS,
-};
-
-/* Reset mock data (for use before each test) */
-static void reset_common_data(enum reset_type t)
-{
-	TEST_SUCC(vb2api_init(workbuf, sizeof(workbuf), &ctx),
-		  "vb2api_init failed");
-
-	memset(&gbb, 0, sizeof(gbb));
-
-	vb2_nv_init(ctx);
-
-	sd = vb2_get_sd(ctx);
-	sd->status |= VB2_SD_STATUS_SECDATA_KERNEL_INIT;
-
-	if (t == FOR_DEVELOPER) {
-		ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
-		sd->flags |= VB2_SD_FLAG_DEV_MODE_ENABLED;
-	}
-
-	ctx->flags |= VB2_CONTEXT_DEV_BOOT_ALLOWED;
-	ctx->flags |= VB2_CONTEXT_DEV_BOOT_EXTERNAL_ALLOWED;
-
-	/* Mock ui_context based on real screens */
-	memset(&mock_ui_context, 0, sizeof(mock_ui_context));
-	mock_ui_context.ctx = ctx;
-	mock_ui_context.state = &mock_state;
-
-	/* For vb2ex_display_ui */
-	memset(mock_displayed, 0, sizeof(mock_displayed));
-	mock_displayed_count = 0;
-	mock_displayed_i = 0;
-
-	/* For vb2ex_get_locale_count */
-	mock_locale_count = 1;
-
-	/* For check_shutdown_request */
-	if (t == FOR_DEVELOPER)
-		mock_calls_until_shutdown = 2000;  /* Larger than 30s */
-	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_total = 0;
-
-	/* For vb2ex_mtime and vb2ex_msleep  */
-	mock_time_ms = mock_time_start_ms;
-
-	/* For vb2ex_beep */
-	memset(mock_beep, 0, sizeof(mock_beep));
-	mock_beep_count = 0;
-	mock_beep_total = 0;
-
-	/* For dev_boot* in 2misc.h */
-	mock_default_boot = VB2_DEV_DEFAULT_BOOT_TARGET_INTERNAL;
-
-	/* For vb2ex_run_altfw */
-	mock_run_altfw_called = 0;
-	mock_altfw_last = -100;
-	mock_altfw_count = 2;
-
-	/* For VbTryLoadKernel */
-	memset(mock_vbtlk_retval, 0, sizeof(mock_vbtlk_retval));
-	memset(mock_vbtlk_expected_flag, 0, sizeof(mock_vbtlk_expected_flag));
-	mock_vbtlk_total = 0;
-
-	/* For vb2_allow_recovery */
-	mock_allow_recovery = t == FOR_MANUAL_RECOVERY;
-
-	/* For vb2ex_physical_presence_pressed */
-	memset(mock_pp_pressed, 0, sizeof(mock_pp_pressed));
-	mock_pp_pressed_total = 0;
-
-	/* For vb2_enable_developer_mode */
-	mock_enable_dev_mode = 0;
-
-	/* For vb2ex_prepare_log_screen */
-	mock_snapshot_count = 0;
-	mock_prepare_log_count = 0;
-	mock_log_page_count = 1;
-
-	/* 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);
-	add_mock_pp_pressed(0);
-
-	mock_diag_storage_test_rv = VB2_SUCCESS;
-}
-
-/* Mock functions */
-struct vb2_gbb_header *vb2_get_gbb(struct vb2_context *c)
-{
-	return &gbb;
-}
-
-vb2_error_t vb2ex_display_ui(enum vb2_screen screen,
-			     uint32_t locale_id,
-			     uint32_t selected_item,
-			     uint32_t disabled_item_mask,
-			     uint32_t hidden_item_mask,
-			     int timer_disabled,
-			     uint32_t current_page,
-			     enum vb2_ui_error error_code)
-{
-	struct display_call displayed = (struct display_call){
-		.screen = vb2_get_screen_info(screen),
-		.locale_id = locale_id,
-		.selected_item = selected_item,
-		.disabled_item_mask = disabled_item_mask,
-		.hidden_item_mask = hidden_item_mask,
-		.timer_disabled = timer_disabled,
-		.current_page = current_page,
-		.error_code = error_code,
-	};
-
-	/* Ignore repeated calls with same arguments */
-	if (mock_displayed_count > 0 &&
-	    !memcmp(&mock_displayed[mock_displayed_count - 1], &displayed,
-		    sizeof(struct display_call)))
-		return VB2_SUCCESS;
-
-	VB2_DEBUG("displayed %d: screen=%#x, locale_id=%u, selected_item=%u, "
-		  "disabled_item_mask=%#x, hidden_item_mask=%#x, "
-		  "timer_disabled=%d, current_page=%u, error=%#x\n",
-		  mock_displayed_count, screen, locale_id, selected_item,
-		  disabled_item_mask, hidden_item_mask,
-		  timer_disabled, current_page, error_code);
-
-	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++] = displayed;
-
-	return VB2_SUCCESS;
-}
-
-uint32_t vb2ex_get_locale_count(void)
-{
-	return mock_locale_count;
-}
-
-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;
-}
-
-uint32_t VbExKeyboardRead(void)
-{
-	return VbExKeyboardReadWithFlags(NULL);
-}
-
-uint32_t VbExKeyboardReadWithFlags(uint32_t *key_flags)
-{
-	mock_iters++;
-	if (mock_iters < mock_key_total) {
-		if (key_flags != NULL) {
-			if (mock_key_trusted[mock_iters])
-				*key_flags = VB_KEY_FLAG_TRUSTED_KEYBOARD;
-			else
-				*key_flags = 0;
-		}
-		return mock_key[mock_iters];
-	}
-
-	return 0;
-}
-
-uint32_t vb2ex_mtime(void)
-{
-	return mock_time_ms;
-}
-
-void vb2ex_msleep(uint32_t msec)
-{
-	mock_time_ms += msec;
-}
-
-void vb2ex_beep(uint32_t msec, uint32_t frequency)
-{
-	struct beep_call *beep;
-	uint32_t cur_time = mock_time_ms - mock_time_start_ms;
-
-	VB2_DEBUG("beep %d: msec = %d, frequency = %d at %d msec\n",
-		  mock_beep_count, msec, frequency, cur_time);
-
-	if (mock_beep_total > 0) {
-		TEST_TRUE(mock_beep_count < mock_beep_total,
-			  "  too many beep calls!");
-
-		beep = &mock_beep[mock_beep_count];
-
-		VB2_DEBUG("beep expected: msec = %d, frequency = %d, "
-			  "at %d msec\n",
-			  beep->msec, beep->frequency, beep->time_expected);
-
-		TEST_EQ(msec, beep->msec, "  beep duration");
-		TEST_EQ(frequency, beep->frequency, "  beep frequency");
-		check_time(cur_time, beep->time_expected,
-			   "  beep started after expected time");
-	}
-
-	mock_time_ms += msec;
-	mock_beep_count++;
-}
-
-enum vb2_dev_default_boot_target vb2api_get_dev_default_boot_target(
-	struct vb2_context *c)
-{
-	return mock_default_boot;
-}
-
-vb2_error_t vb2ex_run_altfw(uint32_t altfw_id)
-{
-	mock_run_altfw_called++;
-	mock_altfw_last = altfw_id;
-
-	return VB2_SUCCESS;
-}
-
-uint32_t vb2ex_get_altfw_count(void)
-{
-	return mock_altfw_count;
-}
-
-vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t disk_flags)
-{
-	int i = mock_iters;
-
-	/* Return last entry if called too many times */
-	if (i >= mock_vbtlk_total)
-		i = mock_vbtlk_total - 1;
-
-	TEST_EQ(mock_vbtlk_expected_flag[i], disk_flags,
-		"  unexpected disk_flags");
-
-	return mock_vbtlk_retval[i];
-}
-
-int vb2api_allow_recovery(struct vb2_context *c)
-{
-	return mock_allow_recovery;
-}
-
-int vb2ex_physical_presence_pressed(void)
-{
-	if (mock_iters >= mock_pp_pressed_total)
-		return 0;
-
-	return mock_pp_pressed[mock_iters];
-}
-
-vb2_error_t vb2api_enable_developer_mode(struct vb2_context *c)
-{
-	mock_enable_dev_mode = 1;
-	return VB2_SUCCESS;
-}
-
-const char *vb2ex_get_debug_info(struct vb2_context *c)
-{
-	return "mocked debug info";
-}
-
-const char *vb2ex_get_firmware_log(int reset)
-{
-	static char mock_firmware_log_buf[MOCK_PREPARE_LOG_SIZE];
-	if (reset)
-		mock_snapshot_count++;
-	snprintf(mock_firmware_log_buf, MOCK_PREPARE_LOG_SIZE,
-		 "%d", mock_snapshot_count);
-	return mock_firmware_log_buf;
-}
-
-uint32_t vb2ex_prepare_log_screen(enum vb2_screen screen, uint32_t locale_id,
-				  const char *str)
-{
-	if (mock_prepare_log_count < ARRAY_SIZE(mock_prepare_log))
-		strncpy(mock_prepare_log[mock_prepare_log_count],
-			str, MOCK_PREPARE_LOG_SIZE);
-	mock_prepare_log_count++;
-
-	return mock_log_page_count;
-}
-
-vb2_error_t vb2ex_diag_get_storage_test_log(const char **log)
-{
-	*log = "mock";
-	return mock_diag_storage_test_rv;
-}
-
-/* Tests */
-static void diagnostics_screen_tests(void)
-{
-	VB2_DEBUG("Testing diagnostic screens...\n");
-
-	/* Diagnostics screen: disabled and hidden item mask */
-	reset_common_data(FOR_DIAGNOSTICS);
-	TEST_EQ(vb2_diagnostic_menu(ctx), VB2_REQUEST_SHUTDOWN,
-		"diagnostic screen: no disabled or hidden item");
-	DISPLAYED_EQ("diagnostic menu", VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE,
-		     MOCK_IGNORE, 0x0, 0x0, MOCK_IGNORE);
-
-	/* Diagnostics screen */
-	reset_common_data(FOR_DIAGNOSTICS);
-
-	/* #0: Language menu */
-	add_mock_keypress(VB_KEY_UP);
-	add_mock_keypress(VB_KEY_ENTER);
-	add_mock_keypress(VB_KEY_ESC);
-	/* #1: Storage health screen */
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_ENTER);
-	add_mock_keypress(VB_KEY_ESC);
-	/* #2: Short storage self-test screen */
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_ENTER);
-	add_mock_keypress(VB_KEY_ESC);
-	/* #3: Extended storage self-test screen */
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_ENTER);
-	add_mock_keypress(VB_KEY_ESC);
-	/* #4: Quick memory test screen */
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_ENTER);
-	add_mock_keypress(VB_KEY_ESC);
-	/* #5: Full memory test screen */
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_ENTER);
-	add_mock_keypress(VB_KEY_ESC);
-	/* #6: Power off (End of menu) */
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_ENTER);
-	mock_calls_until_shutdown = -1;
-	TEST_EQ(vb2_diagnostic_menu(ctx), VB2_REQUEST_SHUTDOWN,
-		"diagnostic screen");
-
-	DISPLAYED_EQ("default on first button of menu", VB2_SCREEN_DIAGNOSTICS,
-		     MOCK_IGNORE, 1, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	/* #0: Language menu */
-	DISPLAYED_EQ("language selection", VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE,
-		     0, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	DISPLAYED_EQ("#0: language menu", VB2_SCREEN_LANGUAGE_SELECT,
-		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE,
-		     MOCK_IGNORE);
-	DISPLAYED_PASS();
-	/* #1: Storage health screen */
-	DISPLAYED_EQ("storage health button", VB2_SCREEN_DIAGNOSTICS,
-		     MOCK_IGNORE, 1, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	DISPLAYED_EQ("#1: storage screen",
-		     VB2_SCREEN_DIAGNOSTICS_STORAGE_HEALTH, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	DISPLAYED_PASS();
-	/* #2: Short storage self-test screen */
-	DISPLAYED_EQ("short storage self-test button", VB2_SCREEN_DIAGNOSTICS,
-		     MOCK_IGNORE, 2, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	DISPLAYED_EQ("#2: short storage self-test screen",
-		     VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST_SHORT, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	DISPLAYED_PASS();
-	/* #3: Extended storage self-test screen */
-	DISPLAYED_EQ("extended storage self-test button",
-		     VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE, 3, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE);
-	DISPLAYED_EQ("#3: extended storage self-test screen",
-		     VB2_SCREEN_DIAGNOSTICS_STORAGE_TEST_EXTENDED, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	DISPLAYED_PASS();
-	/* #4: Quick memory test screen */
-	DISPLAYED_EQ("quick memory test button", VB2_SCREEN_DIAGNOSTICS,
-		     MOCK_IGNORE, 4, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	DISPLAYED_EQ("#4: quick memory test screen",
-		     VB2_SCREEN_DIAGNOSTICS_MEMORY_QUICK, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	DISPLAYED_PASS();
-	/* #5: Full memory test screen */
-	DISPLAYED_EQ("full memory test button", VB2_SCREEN_DIAGNOSTICS,
-		     MOCK_IGNORE, 5, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	DISPLAYED_EQ("#5: full memory test screen",
-		     VB2_SCREEN_DIAGNOSTICS_MEMORY_FULL, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	DISPLAYED_PASS();
-	/* #6: Power of (End of menu) */
-	DISPLAYED_EQ("power off", VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE, 6,
-		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
-	DISPLAYED_NO_EXTRA();
-
-	/* Diagnostics screen: no nvme */
-	reset_common_data(FOR_DIAGNOSTICS);
-	/* Non-nvme storage returns UNIMPLEMENTED. */
-	mock_diag_storage_test_rv = VB2_ERROR_EX_UNIMPLEMENTED;
-	TEST_EQ(vb2_diagnostic_menu(ctx), VB2_REQUEST_SHUTDOWN,
-		"diagnostic screen: check disabled item");
-	DISPLAYED_EQ("diagnostic menu: self-test disabled",
-		     VB2_SCREEN_DIAGNOSTICS, MOCK_IGNORE, MOCK_IGNORE, 0xc, 0x0,
-		     MOCK_IGNORE);
-
-	VB2_DEBUG("...done.\n");
-}
-
-int main(void)
-{
-	/* Screen displayed */
-	diagnostics_screen_tests();
-
-	return gTestSuccess ? 0 : 255;
-}
diff --git a/tests/vboot_api_kernel4_tests.c b/tests/vboot_api_kernel4_tests.c
index 14c6ea34..f122eb38 100644
--- a/tests/vboot_api_kernel4_tests.c
+++ b/tests/vboot_api_kernel4_tests.c
@@ -11,7 +11,6 @@
 #include "2nvstorage.h"
 #include "2secdata.h"
 #include "2sysincludes.h"
-#include "2ui.h"
 #include "host_common.h"
 #include "load_kernel_fw.h"
 #include "test_common.h"
@@ -165,7 +164,7 @@ vb2_error_t vb2ex_broken_screen_ui(struct vb2_context *c)
 	return vbboot_retval;
 }
 
-vb2_error_t vb2_diagnostic_menu(struct vb2_context *c)
+vb2_error_t vb2ex_diagnostic_ui(struct vb2_context *c)
 {
 	if (vbboot_retval == -5)
 		return VB2_ERROR_MOCK;