minidiag/ui: Prevent calling update action if memory test is finished

Skip calling diagnostics_memory_update_screen() action if memory test is
finished to reduce delays in memory diagnostics screens.

BUG=b:168776970, b:171852478
BRANCH=none
TEST=Build locally, boot recovery, select 'run diagnostics',
     enter memory test (quick) screen, wait until test is finished,
     and observe that the delay reduces while pressing keyboard

Signed-off-by: Hsuan Ting Chen <roccochen@chromium.org>
Change-Id: I4f8dbadc68989112bce418119ff7405ec150a7df
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2423089
diff --git a/firmware/2lib/2ui_screens.c b/firmware/2lib/2ui_screens.c
index 1b4c5e2..878b4ed 100644
--- a/firmware/2lib/2ui_screens.c
+++ b/firmware/2lib/2ui_screens.c
@@ -1123,6 +1123,11 @@
 						    int reset)
 {
 	const char *log_string = NULL;
+
+	/* Early return if the memory test is done. */
+	if (ui->state->test_finished)
+		return VB2_REQUEST_UI_CONTINUE;
+
 	vb2_error_t rv = op(reset, &log_string);
 	if ((rv && rv != VB2_ERROR_EX_DIAG_TEST_RUNNING) || !log_string) {
 		VB2_DEBUG("ERROR: Failed to retrieve memory test status\n");
@@ -1158,6 +1163,7 @@
 			    DIAGNOSTICS_MEMORY_ITEM_CANCEL);
 		if (ui->state->selected_item == DIAGNOSTICS_MEMORY_ITEM_CANCEL)
 			ui->state->selected_item = DIAGNOSTICS_MEMORY_ITEM_BACK;
+		ui->state->test_finished = 1;
 	}
 
 	return VB2_REQUEST_UI_CONTINUE;
diff --git a/firmware/2lib/include/2ui.h b/firmware/2lib/include/2ui.h
index 44e24fd..77c5acf 100644
--- a/firmware/2lib/include/2ui.h
+++ b/firmware/2lib/include/2ui.h
@@ -78,6 +78,9 @@
 	uint32_t page_count;
 	uint32_t current_page;
 
+	/* For memory check screen. */
+	int test_finished;  /* Do not update screen if the content is done */
+
 	struct vb2_screen_state *prev;
 };