1 files changed, 55 insertions, 9 deletions

diff --git a/test/boot/bootflow.c b/test/boot/bootflow.c
index f640db8a24..b97c566f00 100644
--- a/test/boot/bootflow.c
+++ b/test/boot/bootflow.c
@@ -511,19 +511,27 @@ BOOTSTD_TEST(bootflow_cmd_boot, UT_TESTF_DM | UT_TESTF_SCAN_FDT);
 /**
  * prep_mmc_bootdev() - Set up an mmc bootdev so we can access other distros
  *
+ * After calling this function, set std->bootdev_order to *@old_orderp to
+ * restore normal operation of bootstd (i.e. with the original bootdev order)
+ *
  * @uts: Unit test state
- * @mmc_dev: MMC device to use, e.g. "mmc4"
+ * @mmc_dev: MMC device to use, e.g. "mmc4". Note that this must remain valid
+ *	in the caller until
+ * @bind_cros: true to bind the ChromiumOS bootmeth
+ * @old_orderp: Returns the original bootdev order, which must be restored
  * Returns 0 on success, -ve on failure
  */
 static int prep_mmc_bootdev(struct unit_test_state *uts, const char *mmc_dev,
-			    bool bind_cros)
+			    bool bind_cros, const char ***old_orderp)
 {
-	const char *order[] = {"mmc2", "mmc1", mmc_dev, NULL};
+	static const char *order[] = {"mmc2", "mmc1", NULL, NULL};
 	struct udevice *dev, *bootstd;
 	struct bootstd_priv *std;
 	const char **old_order;
 	ofnode root, node;
 
+	order[2] = mmc_dev;
+
 	/* Enable the mmc4 node since we need a second bootflow */
 	root = oftree_root(oftree_default());
 	node = ofnode_find_subnode(root, mmc_dev);
@@ -546,26 +554,49 @@ static int prep_mmc_bootdev(struct unit_test_state *uts, const char *mmc_dev,
 	std = dev_get_priv(bootstd);
 	old_order = std->bootdev_order;
 	std->bootdev_order = order;
+	*old_orderp = old_order;
+
+	return 0;
+}
+
+/**
+ * scan_mmc_bootdev() - Set up an mmc bootdev so we can access other distros
+ *
+ * @uts: Unit test state
+ * @mmc_dev: MMC device to use, e.g. "mmc4"
+ * @bind_cros: true to bind the ChromiumOS bootmeth
+ * Returns 0 on success, -ve on failure
+ */
+static int scan_mmc_bootdev(struct unit_test_state *uts, const char *mmc_dev,
+			    bool bind_cros)
+{
+	struct bootstd_priv *std;
+	struct udevice *bootstd;
+	const char **old_order;
+
+	ut_assertok(prep_mmc_bootdev(uts, mmc_dev, bind_cros, &old_order));
 
 	console_record_reset_enable();
 	ut_assertok(run_command("bootflow scan", 0));
 	ut_assert_console_end();
 
 	/* Restore the order used by the device tree */
+	ut_assertok(uclass_first_device_err(UCLASS_BOOTSTD, &bootstd));
+	std = dev_get_priv(bootstd);
 	std->bootdev_order = old_order;
 
 	return 0;
 }
 
 /**
- * prep_mmc4_bootdev() - Set up the mmc4 bootdev so we can access a fake Armbian
+ * scan_mmc4_bootdev() - Set up the mmc4 bootdev so we can access a fake Armbian
  *
  * @uts: Unit test state
  * Returns 0 on success, -ve on failure
  */
-static int prep_mmc4_bootdev(struct unit_test_state *uts)
+static int scan_mmc4_bootdev(struct unit_test_state *uts)
 {
-	ut_assertok(prep_mmc_bootdev(uts, "mmc4", false));
+	ut_assertok(scan_mmc_bootdev(uts, "mmc4", false));
 
 	return 0;
 }
@@ -573,9 +604,13 @@ static int prep_mmc4_bootdev(struct unit_test_state *uts)
 /* Check 'bootflow menu' to select a bootflow */
 static int bootflow_cmd_menu(struct unit_test_state *uts)
 {
+	struct bootstd_priv *std;
 	char prev[3];
 
-	ut_assertok(prep_mmc4_bootdev(uts));
+	/* get access to the current bootflow */
+	ut_assertok(bootstd_get_priv(&std));
+
+	ut_assertok(scan_mmc4_bootdev(uts));
 
 	/* Add keypresses to move to and select the second one in the list */
 	prev[0] = CTL_CH('n');
@@ -585,6 +620,17 @@ static int bootflow_cmd_menu(struct unit_test_state *uts)
 
 	ut_assertok(run_command("bootflow menu", 0));
 	ut_assert_nextline("Selected: Armbian");
+	ut_assertnonnull(std->cur_bootflow);
+	ut_assert_console_end();
+
+	/* Check not selecting anything */
+	prev[0] = '\e';
+	prev[1] = '\0';
+	ut_asserteq(1, console_in_puts(prev));
+
+	ut_asserteq(1, run_command("bootflow menu", 0));
+	ut_assertnull(std->cur_bootflow);
+	ut_assert_nextline("Nothing chosen");
 	ut_assert_console_end();
 
 	return 0;
@@ -681,7 +727,7 @@ static int bootflow_menu_theme(struct unit_test_state *uts)
 	ofnode node;
 	int i;
 
-	ut_assertok(prep_mmc4_bootdev(uts));
+	ut_assertok(scan_mmc4_bootdev(uts));
 
 	ut_assertok(bootflow_menu_new(&exp));
 	node = ofnode_path("/bootstd/theme");
@@ -996,7 +1042,7 @@ BOOTSTD_TEST(bootflow_cmdline_special, 0);
 /* Test ChromiumOS bootmeth */
 static int bootflow_cros(struct unit_test_state *uts)
 {
-	ut_assertok(prep_mmc_bootdev(uts, "mmc5", true));
+	ut_assertok(scan_mmc_bootdev(uts, "mmc5", true));
 	ut_assertok(run_command("bootflow list", 0));
 
 	ut_assert_nextlinen("Showing all");