vboot/ui: Add tests for to_dev transition flowfirmware-endeavour-13259.B

Unit tests paired with CL:2168072.

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

Cq-Depend: chromium:2214457
Signed-off-by: Hsuan Ting Chen <roccochen@chromium.org>
Change-Id: Ic6b109c514e1582d7eb29040135aeaa884b243be
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2198274
Reviewed-by: Yu-Ping Wu <yupingso@chromium.org>
Commit-Queue: Yu-Ping Wu <yupingso@chromium.org>

3 files changed, 320 insertions, 32 deletions

diff --git a/tests/vb2_ui_action_tests.c b/tests/vb2_ui_action_tests.c
index ab273da0..ddb9ee38 100644
--- a/tests/vb2_ui_action_tests.c
+++ b/tests/vb2_ui_action_tests.c
@@ -24,7 +24,8 @@
 #define MOCK_SCREEN_TARGET0 0xef20
 #define MOCK_SCREEN_TARGET1 0xef21
 #define MOCK_SCREEN_TARGET2 0xef22
-#define MOCK_SCREEN_TARGET3 0xef23
+#define MOCK_SCREEN_ACTION 0xef30
+#define MOCK_SCREEN_ALL_ACTION 0xef32
 
 /* Mock data */
 struct display_call {
@@ -37,6 +38,7 @@ struct display_call {
 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 int mock_calls_until_shutdown;
@@ -60,9 +62,10 @@ static uint32_t mock_vbtlk_expected_flag;
 
 /* Mock actions */
 static uint32_t mock_action_called;
+static uint32_t mock_action_countdown_limit;
 static vb2_error_t mock_action_countdown(struct vb2_ui_context *ui)
 {
-	if (++mock_action_called >= 10)
+	if (++mock_action_called >= mock_action_countdown_limit)
 		return VB2_SUCCESS;
 	return VB2_REQUEST_UI_CONTINUE;
 }
@@ -72,6 +75,34 @@ static vb2_error_t mock_action_change_screen(struct vb2_ui_context *ui)
 	return vb2_ui_change_screen(ui, MOCK_SCREEN_BASE);
 }
 
+static vb2_error_t mock_action_base(struct vb2_ui_context *ui)
+{
+	mock_action_called++;
+	return VB2_SUCCESS;
+}
+
+static int mock_action_flags;
+static vb2_error_t mock_action_flag0(struct vb2_ui_context *ui)
+{
+	if ((1 << 0) & mock_action_flags)
+		return VB2_SUCCESS;
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
+static vb2_error_t mock_action_flag1(struct vb2_ui_context *ui)
+{
+	if ((1 << 1) & mock_action_flags)
+		return VB2_SUCCESS;
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
+static vb2_error_t mock_action_flag2(struct vb2_ui_context *ui)
+{
+	if ((1 << 2) & mock_action_flags)
+		return VB2_SUCCESS;
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
 /* Mock screens */
 struct vb2_screen_info mock_screen_temp;
 const struct vb2_menu_item mock_empty_menu[] = {};
@@ -89,28 +120,28 @@ const struct vb2_screen_info mock_screen_base = {
 };
 const struct vb2_menu_item mock_screen_menu_items[] = {
 	{
-		.text = "option 0",
+		.text = "item 0",
 		.target = MOCK_SCREEN_TARGET0,
 	},
 	{
-		.text = "option 1",
+		.text = "item 1",
 		.target = MOCK_SCREEN_TARGET1,
 	},
 	{
-		.text = "option 2",
+		.text = "item 2",
 		.target = MOCK_SCREEN_TARGET2,
 	},
 	{
-		.text = "option 3",
-		.target = MOCK_SCREEN_TARGET3,
+		.text = "item 3",
+		.action = mock_action_base,
 	},
 	{
-		.text = "option 4 (no target)",
+		.text = "item 4 (no target)",
 	},
 };
 const struct vb2_screen_info mock_screen_menu = {
 	.id = MOCK_SCREEN_MENU,
-	.name = "mock_screen_menu: screen with 5 options",
+	.name = "mock_screen_menu: screen with 5 items",
 	.num_items = ARRAY_SIZE(mock_screen_menu_items),
 	.items = mock_screen_menu_items,
 };
@@ -132,12 +163,26 @@ const struct vb2_screen_info 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",
+const struct vb2_screen_info mock_screen_action = {
+	.id = MOCK_SCREEN_ACTION,
+	.name = "mock_screen_action",
+	.action = mock_action_countdown,
 	.num_items = ARRAY_SIZE(mock_empty_menu),
 	.items = mock_empty_menu,
 };
+const struct vb2_menu_item mock_screen_all_action_items[] = {
+	{
+		.text = "all_action_screen_item",
+		.action = mock_action_flag1,
+	},
+};
+const struct vb2_screen_info mock_screen_all_action = {
+	.id = MOCK_SCREEN_ALL_ACTION,
+	.name = "mock_screen_all_action",
+	.action = mock_action_flag0,
+	.num_items = ARRAY_SIZE(mock_screen_all_action_items),
+	.items = mock_screen_all_action_items,
+};
 
 static void screen_state_eq(const struct vb2_screen_state *state,
 			    enum vb2_screen screen,
@@ -239,6 +284,8 @@ static void reset_common_data(void)
 
 	vb2_nv_init(ctx);
 
+	sd = vb2_get_sd(ctx);
+
 	/* For check_shutdown_request */
 	mock_calls_until_shutdown = 10;
 
@@ -253,6 +300,7 @@ static void reset_common_data(void)
 	/* Mock ui_context based on mock screens */
 	memset(&mock_ui_context, 0, sizeof(mock_ui_context));
 	mock_ui_context.ctx = ctx;
+	mock_ui_context.state.screen = &mock_screen_temp;
 	mock_state = &mock_ui_context.state;
 
 	/* For vb2ex_display_ui */
@@ -268,6 +316,8 @@ static void reset_common_data(void)
 
 	/* For mock actions */
 	mock_action_called = 0;
+	mock_action_countdown_limit = 1;
+	mock_action_flags = 0;
 
 	/* For chagen_screen and vb2_get_screen_info */
 	mock_get_screen_info_called = 0;
@@ -311,8 +361,10 @@ const struct vb2_screen_info *vb2_get_screen_info(enum vb2_screen screen)
 		return &mock_screen_target1;
 	case MOCK_SCREEN_TARGET2:
 		return &mock_screen_target2;
-	case MOCK_SCREEN_TARGET3:
-		return &mock_screen_target3;
+	case MOCK_SCREEN_ACTION:
+		return &mock_screen_action;
+	case MOCK_SCREEN_ALL_ACTION:
+		return &mock_screen_all_action;
 	case MOCK_NO_SCREEN:
 		return NULL;
 	default:
@@ -491,9 +543,6 @@ static void menu_next_tests(void)
 
 static void menu_select_tests(void)
 {
-	int i, target_id;
-	char test_name[256];
-
 	VB2_DEBUG("Testing menu_select...\n");
 
 	/* select action with no item screen */
@@ -505,27 +554,32 @@ static void menu_select_tests(void)
 		"vb2_ui_menu_select 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(vb2_ui_menu_select(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE, test_name);
-		screen_state_eq(mock_state, target_id, 0, MOCK_IGNORE);
-	}
+	/* Try to select an item with a target (item 2) */
+	reset_common_data();
+	mock_state->screen = &mock_screen_menu;
+	mock_state->selected_item = 2;
+	mock_ui_context.key = VB_KEY_ENTER;
+	TEST_EQ(vb2_ui_menu_select(&mock_ui_context),
+		VB2_REQUEST_UI_CONTINUE, "select an item with a target");
+	screen_state_eq(mock_state, MOCK_SCREEN_TARGET2, 0, MOCK_IGNORE);
+
+	/* Try to select an item with an action (item 3) */
+	reset_common_data();
+	mock_state->screen = &mock_screen_menu;
+	mock_state->selected_item = 3;
+	mock_ui_context.key = VB_KEY_ENTER;
+	TEST_EQ(vb2_ui_menu_select(&mock_ui_context),
+		VB2_SUCCESS, "select an item with an action");
+	TEST_EQ(mock_action_called, 1, "  action called once");
 
-	/* Try to select no target item (target 4) */
+	/* Try to select an item with neither targets nor actions (item 4) */
 	reset_common_data();
 	mock_state->screen = &mock_screen_menu;
 	mock_state->selected_item = 4;
 	mock_ui_context.key = VB_KEY_ENTER;
 	TEST_EQ(vb2_ui_menu_select(&mock_ui_context),
 		VB2_REQUEST_UI_CONTINUE,
-		"select no target");
+		"select an item with neither targets nor actions");
 	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 4, MOCK_IGNORE);
 
 	/* Ignore power button short press when not DETACHABLE */
@@ -606,6 +660,13 @@ static void manual_recovery_action_tests(void)
 
 static void ui_loop_tests(void)
 {
+	int i;
+	const char *action_interfere_test_names[] = {
+		"hook all actions: screen action return SUCCESS",
+		"hook all actions: target action hooked return SUCCESS",
+		"hook all actions: global action return SUCCESS",
+	};
+
 	VB2_DEBUG("Testing ui_loop...\n");
 
 	/* Die if no root screen */
@@ -623,9 +684,18 @@ static void ui_loop_tests(void)
 		     MOCK_IGNORE, MOCK_IGNORE);
 	displayed_no_extra();
 
+	/* Screen action */
+	reset_common_data();
+	mock_calls_until_shutdown = -1;
+	mock_action_countdown_limit = 10;
+	TEST_EQ(ui_loop(ctx, MOCK_SCREEN_ACTION, NULL),
+		VB2_SUCCESS, "screen action");
+	TEST_EQ(mock_action_called, 10, "  action called");
+
 	/* Global action */
 	reset_common_data();
 	mock_calls_until_shutdown = -1;
+	mock_action_countdown_limit = 10;
 	TEST_EQ(ui_loop(ctx, VB2_SCREEN_BLANK, mock_action_countdown),
 		VB2_SUCCESS, "global action");
 	TEST_EQ(mock_action_called, 10, "  action called");
@@ -639,6 +709,21 @@ static void ui_loop_tests(void)
 	displayed_eq("change to mock_screen_base", MOCK_IGNORE, MOCK_IGNORE,
 		     MOCK_IGNORE, MOCK_IGNORE);
 
+	/*
+	 * Hook all actions, and receive SUCCESS from actions one by one
+	 * Action #0: screen action
+	 * Action #1: item target action
+	 * Action #2: global action
+	 */
+	for (i = 0; i <= 2; i++) {
+		reset_common_data();
+		add_mock_keypress(VB_KEY_ENTER);
+		mock_calls_until_shutdown = -1;
+		mock_action_flags |= (1 << i);
+		TEST_EQ(ui_loop(ctx, MOCK_SCREEN_ALL_ACTION, mock_action_flag2),
+			VB2_SUCCESS, action_interfere_test_names[i]);
+	}
+
 	/* KEY_UP, KEY_DOWN, and KEY_ENTER navigation */
 	reset_common_data();
 	add_mock_keypress(VB_KEY_UP);  /* (blocked) */
diff --git a/tests/vb2_ui_tests.c b/tests/vb2_ui_tests.c
index 325e8423..4066e5b8 100644
--- a/tests/vb2_ui_tests.c
+++ b/tests/vb2_ui_tests.c
@@ -66,6 +66,14 @@ 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;
+
 static void add_mock_key(uint32_t press, int trusted)
 {
 	if (mock_key_total >= ARRAY_SIZE(mock_key) ||
@@ -97,6 +105,16 @@ static void add_mock_vbtlk(vb2_error_t retval, uint32_t get_info_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,
@@ -210,6 +228,16 @@ static void reset_common_data(enum reset_type t)
 	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;
+
 	/* Avoid Iteration #0 */
 	add_mock_keypress(0);
 	if (t == FOR_MANUAL_RECOVERY)
@@ -217,6 +245,7 @@ static void reset_common_data(enum reset_type t)
 			       VB_DISK_FLAG_REMOVABLE);
 	else
 		add_mock_vbtlk(VB2_ERROR_MOCK, 0);
+	add_mock_pp_pressed(0);
 }
 
 /* Mock functions */
@@ -341,6 +370,24 @@ vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t get_info_flags)
 	return mock_vbtlk_retval[i];
 }
 
+int vb2_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];
+}
+
+void vb2_enable_developer_mode(struct vb2_context *c)
+{
+	mock_enable_dev_mode = 1;
+}
+
 /* Tests */
 static void developer_tests(void)
 {
@@ -513,6 +560,146 @@ static void manual_recovery_tests(void)
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	displayed_no_extra();
 
+	/* Ctrl+D = to_dev; space = cancel */
+	reset_common_data(FOR_MANUAL_RECOVERY);
+	add_mock_key(VB_KEY_CTRL('D'), 1);
+	add_mock_keypress(' ');
+	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
+		"ctrl+D = to_dev; space = cancel");
+	TEST_EQ(mock_enable_dev_mode, 0, "  dev mode not enabled");
+	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_eq("to_dev", VB2_SCREEN_RECOVERY_TO_DEV,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_no_extra();
+
+	/* Cancel */
+	reset_common_data(FOR_MANUAL_RECOVERY);
+	add_mock_key(VB_KEY_CTRL('D'), 1);
+	if (PHYSICAL_PRESENCE_KEYBOARD)
+		add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_ENTER);
+	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN, "cancel");
+	TEST_EQ(mock_enable_dev_mode, 0, "  dev mode not enabled");
+
+	/* Confirm */
+	reset_common_data(FOR_MANUAL_RECOVERY);
+	add_mock_key(VB_KEY_CTRL('D'), 1);
+	if (PHYSICAL_PRESENCE_KEYBOARD) {
+		add_mock_key(VB_KEY_ENTER, 1);
+	} else {
+		add_mock_pp_pressed(0);
+		add_mock_pp_pressed(1);
+		add_mock_pp_pressed(1);
+		add_mock_pp_pressed(0);
+	}
+	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_REBOOT_EC_TO_RO,
+		"confirm");
+	if (!PHYSICAL_PRESENCE_KEYBOARD)
+		TEST_TRUE(mock_iters >= mock_pp_pressed_total - 1,
+			  "  used up mock_pp_pressed");
+	TEST_EQ(mock_enable_dev_mode, 1, "  dev mode enabled");
+
+	/* Cannot confirm physical presence by untrusted keyboard */
+	if (PHYSICAL_PRESENCE_KEYBOARD) {
+		reset_common_data(FOR_MANUAL_RECOVERY);
+		add_mock_key(VB_KEY_CTRL('D'), 1);
+		add_mock_key(VB_KEY_ENTER, 0);
+		TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
+			"cannot confirm physical presence"
+			" by untrusted keyboard");
+		TEST_EQ(mock_enable_dev_mode, 0, "  dev mode not enabled");
+	}
+
+	/* Cannot enable dev mode if already enabled */
+	reset_common_data(FOR_MANUAL_RECOVERY);
+	sd->flags |= VB2_SD_FLAG_DEV_MODE_ENABLED;
+	add_mock_key(VB_KEY_CTRL('D'), 1);
+	if (PHYSICAL_PRESENCE_KEYBOARD) {
+		add_mock_key(VB_KEY_ENTER, 1);
+	} else {
+		add_mock_pp_pressed(0);
+		add_mock_pp_pressed(1);
+		add_mock_pp_pressed(0);
+	}
+	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
+		"cannot enable dev mode if already enabled");
+	TEST_EQ(mock_enable_dev_mode, 0, "  dev mode already on");
+
+	/* Physical presence button tests */
+	if (!PHYSICAL_PRESENCE_KEYBOARD) {
+		/* Physical presence button stuck? */
+		reset_common_data(FOR_MANUAL_RECOVERY);
+		add_mock_key(VB_KEY_CTRL('D'), 1);
+		add_mock_pp_pressed(1);  /* Hold since boot */
+		add_mock_pp_pressed(0);
+		TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
+			"physical presence button stuck?");
+		TEST_EQ(mock_enable_dev_mode, 0, "  dev mode not enabled");
+		displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+			     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+		displayed_no_extra();
+
+		/* Button stuck, enter to_dev again */
+		reset_common_data(FOR_MANUAL_RECOVERY);
+		add_mock_key(VB_KEY_CTRL('D'), 1);
+		add_mock_key(VB_KEY_CTRL('D'), 1);
+		add_mock_pp_pressed(1);  /* Hold since boot */
+		add_mock_pp_pressed(0);
+		add_mock_pp_pressed(1);  /* Press again */
+		add_mock_pp_pressed(0);
+		TEST_EQ(vb2_manual_recovery_menu(ctx),
+			VB2_REQUEST_REBOOT_EC_TO_RO,
+			"button stuck, enter to_dev again");
+		TEST_TRUE(mock_iters >= mock_pp_pressed_total - 1,
+			  "  used up mock_pp_pressed");
+		TEST_EQ(mock_enable_dev_mode, 1, "  dev mode enabled");
+		displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+			     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+		displayed_eq("to_dev", VB2_SCREEN_RECOVERY_TO_DEV,
+			     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+		displayed_no_extra();
+
+		/* Cancel with holding pp button, enter again */
+		reset_common_data(FOR_MANUAL_RECOVERY);
+		/* Enter to_dev */
+		add_mock_key(VB_KEY_CTRL('D'), 1);
+		add_mock_pp_pressed(0);
+		/* Press pp button */
+		add_mock_keypress(0);
+		add_mock_pp_pressed(1);
+		/* Space = back */
+		add_mock_keypress(' ');
+		add_mock_pp_pressed(1);
+		/* Wait */
+		add_mock_keypress(0);
+		add_mock_pp_pressed(0);
+		/* Enter to_dev again */
+		add_mock_key(VB_KEY_CTRL('D'), 1);
+		add_mock_pp_pressed(0);
+		/* Press pp button again */
+		add_mock_pp_pressed(1);
+		/* Release */
+		add_mock_pp_pressed(0);
+		TEST_EQ(vb2_manual_recovery_menu(ctx),
+			VB2_REQUEST_REBOOT_EC_TO_RO,
+			"cancel with holding pp button, enter again");
+		TEST_TRUE(mock_iters >= mock_pp_pressed_total - 1,
+			  "  used up mock_pp_pressed");
+		TEST_EQ(mock_enable_dev_mode, 1, "  dev mode enabled");
+		displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+			     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+		displayed_eq("to_dev", VB2_SCREEN_RECOVERY_TO_DEV,
+			     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+		displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+			     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+		displayed_eq("to_dev", VB2_SCREEN_RECOVERY_TO_DEV,
+			     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+		displayed_no_extra();
+	}
+
 	VB2_DEBUG("...done.\n");
 }
 
diff --git a/tests/vb2_ui_utility_tests.c b/tests/vb2_ui_utility_tests.c
index 6fdaffc1..257e69d9 100644
--- a/tests/vb2_ui_utility_tests.c
+++ b/tests/vb2_ui_utility_tests.c
@@ -34,6 +34,14 @@ static int mock_shutdown_request;
 static struct vb2_ui_context mock_ui_context;
 static struct vb2_screen_state *mock_state;
 
+/* Mock actions */
+static uint32_t mock_action_called;
+static vb2_error_t mock_action_base(struct vb2_ui_context *ui)
+{
+	mock_action_called++;
+	return VB2_SUCCESS;
+}
+
 /* Mock screens */
 const struct vb2_menu_item mock_empty_menu[] = {};
 struct vb2_screen_info mock_screen_blank = {
@@ -114,6 +122,9 @@ static void reset_common_data(void)
 	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;
+
+	/* For mock actions */
+	mock_action_called = 0;
 }
 
 /* Mock functions */
@@ -281,7 +292,12 @@ 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);
 
-	/* TODO: Change to screen with init */
+	/* Change to screen with init */
+	reset_common_data();
+	mock_screen_base.init = mock_action_base;
+	TEST_EQ(vb2_ui_change_screen(&mock_ui_context, MOCK_SCREEN_BASE),
+		VB2_SUCCESS, "change to screen with init");
+	TEST_EQ(mock_action_called, 1, "  action called once");
 
 	VB2_DEBUG("...done.\n");
 }