[Concept,03/11] bootstage: Add some more tests

diff --git a/common/bootstage.c b/common/bootstage.c
index 1fee2edeef3..e0298991fa9 100644
--- a/common/bootstage.c
+++ b/common/bootstage.c
@@ -212,6 +212,43 @@  uint32_t bootstage_accum(enum bootstage_id id)
 	return duration;
 }
 
+uint bootstage_get_rec_count(void)
+{
+	struct bootstage_data *data = gd->bootstage;
+
+	if (!data)
+		return 0;
+
+	return data->rec_count;
+}
+
+const struct bootstage_record *bootstage_get_rec(uint index)
+{
+	struct bootstage_data *data = gd->bootstage;
+
+	if (!data || index >= data->rec_count)
+		return NULL;
+
+	return &data->record[index];
+}
+
+void bootstage_set_rec_count(uint count)
+{
+	struct bootstage_data *data = gd->bootstage;
+	uint i;
+
+	if (!data || count > RECORD_COUNT)
+		return;
+
+	/* Clear any records beyond the new count */
+	for (i = count; i < data->rec_count; i++) {
+		data->record[i].time_us = 0;
+		data->record[i].start_us = 0;
+	}
+
+	data->rec_count = count;
+}
+
 /**
  * Get a record name as a printable string
  *
diff --git a/include/bootstage.h b/include/bootstage.h
index 98bd9024da8..adc7b8f0a35 100644
--- a/include/bootstage.h
+++ b/include/bootstage.h
@@ -354,6 +354,30 @@  uint32_t bootstage_start(enum bootstage_id id, const char *name);
  */
 uint32_t bootstage_accum(enum bootstage_id id);
 
+/**
+ * bootstage_get_rec_count() - Get the number of bootstage records
+ *
+ * Return: number of bootstage records
+ */
+uint bootstage_get_rec_count(void);
+
+/**
+ * bootstage_get_rec() - Get a bootstage record by index
+ *
+ * @index: Index of the record to retrieve (numbered from 0)
+ * Return: pointer to the record, or NULL if @index is out of range
+ */
+const struct bootstage_record *bootstage_get_rec(uint index);
+
+/**
+ * bootstage_set_rec_count() - Set the number of bootstage records
+ *
+ * This can be used to restore the record count after testing
+ *
+ * @count: New record count (must be <= RECORD_COUNT)
+ */
+void bootstage_set_rec_count(uint count);
+
 /* Print a report about boot time */
 void bootstage_report(void);
 
diff --git a/test/cmd/Makefile b/test/cmd/Makefile
index 3fc07f0cacf..9cd8ea3aaf0 100644
--- a/test/cmd/Makefile
+++ b/test/cmd/Makefile
@@ -15,6 +15,7 @@  obj-$(CONFIG_X86) += cpuid.o msr.o
 obj-$(CONFIG_CMD_ADDR_FIND) += addr_find.o
 obj-$(CONFIG_CMD_ADDRMAP) += addrmap.o
 obj-$(CONFIG_CMD_BDI) += bdinfo.o
+obj-$(CONFIG_CMD_BOOTSTAGE) += bootstage.o
 obj-$(CONFIG_CMD_CHID) += chid.o
 obj-$(CONFIG_COREBOOT_SYSINFO) += coreboot.o
 obj-$(CONFIG_CMD_FDT) += fdt.o
diff --git a/test/cmd/bootstage.c b/test/cmd/bootstage.c
new file mode 100644
index 00000000000..dee4a0671fa
--- /dev/null
+++ b/test/cmd/bootstage.c
@@ -0,0 +1,32 @@ 
+// SPDX-License-Identifier: GPL-2.0+
+/*
+ * Tests for bootstage command
+ *
+ * Copyright 2025 Canonical Ltd
+ */
+
+#include <bootstage.h>
+#include <test/cmd.h>
+#include <test/ut.h>
+
+static int cmd_bootstage_report(struct unit_test_state *uts)
+{
+	uint count;
+
+	/* Get the current record count */
+	count = bootstage_get_rec_count();
+	ut_assert(count > 0);
+
+	/* Test the bootstage report command runs successfully */
+	ut_assertok(run_command("bootstage report", 0));
+
+	/* Verify the report contains expected headers and stages */
+	ut_assert_nextline("Timer summary in microseconds (%u records):",
+			   count);
+	ut_assert_nextline("       Mark    Elapsed  Stage");
+	ut_assert_nextline("          0          0  reset");
+	ut_assert_skip_to_line("Accumulated time:");
+
+	return 0;
+}
+CMD_TEST(cmd_bootstage_report, UTF_CONSOLE);
diff --git a/test/common/Makefile b/test/common/Makefile
index 7b5927b5a44..a5df946396a 100644
--- a/test/common/Makefile
+++ b/test/common/Makefile
@@ -7,6 +7,8 @@  obj-$(CONFIG_$(PHASE_)CMDLINE) += bloblist.o
 endif
 endif
 
+obj-$(CONFIG_BOOTSTAGE) += bootstage.o
+
 obj-$(CONFIG_CONSOLE_PAGER) += console.o
 obj-$(CONFIG_CYCLIC) += cyclic.o
 obj-$(CONFIG_EVENT_DYNAMIC) += event.o
diff --git a/test/common/bootstage.c b/test/common/bootstage.c
new file mode 100644
index 00000000000..6e97c3e3a72
--- /dev/null
+++ b/test/common/bootstage.c
@@ -0,0 +1,199 @@ 
+// SPDX-License-Identifier: GPL-2.0+
+/*
+ * Tests for bootstage API
+ *
+ * Copyright 2025 Canonical Ltd
+ */
+
+#include <bootstage.h>
+#include <linux/delay.h>
+#include <test/common.h>
+#include <test/test.h>
+#include <test/ut.h>
+
+/* Test bootstage_mark_name() */
+static int test_bootstage_mark(struct unit_test_state *uts)
+{
+	const struct bootstage_record *rec;
+	ulong time;
+	int count;
+
+	/* Get the current count so we know where our record will be */
+	count = bootstage_get_rec_count();
+
+	/* Mark a stage and verify we get a valid timestamp */
+	time = bootstage_mark_name(BOOTSTAGE_ID_USER + 50, "test_stage_mark");
+	ut_assert(time > 0);
+
+	/* Verify the count increased by 1 */
+	ut_asserteq(count + 1, bootstage_get_rec_count());
+
+	/* Check that the record was added correctly */
+	rec = bootstage_get_rec(count);
+	ut_assertnonnull(rec);
+	ut_asserteq(BOOTSTAGE_ID_USER + 50, rec->id);
+	ut_asserteq_str("test_stage_mark", rec->name);
+	ut_asserteq(time, rec->time_us);
+	ut_asserteq(0, rec->flags);
+
+	/* Restore the original count */
+	bootstage_set_rec_count(count);
+
+	return 0;
+}
+COMMON_TEST(test_bootstage_mark, 0);
+
+/* Test bootstage_error_name() */
+static int test_bootstage_error(struct unit_test_state *uts)
+{
+	const struct bootstage_record *rec;
+	ulong time;
+	int count;
+
+	count = bootstage_get_rec_count();
+
+	/* Mark an error stage and verify we get a valid timestamp */
+	time = bootstage_error_name(BOOTSTAGE_ID_USER + 51, "test_error");
+	ut_assert(time > 0);
+
+	/* Check the error record */
+	rec = bootstage_get_rec(count);
+	ut_assertnonnull(rec);
+	ut_asserteq(BOOTSTAGE_ID_USER + 51, rec->id);
+	ut_asserteq_str("test_error", rec->name);
+	ut_asserteq(time, rec->time_us);
+	ut_asserteq(BOOTSTAGEF_ERROR, rec->flags);
+
+	/* Restore the original count */
+	bootstage_set_rec_count(count);
+
+	return 0;
+}
+COMMON_TEST(test_bootstage_error, 0);
+
+/* Test bootstage_start() and bootstage_accum() */
+static int test_bootstage_accum(struct unit_test_state *uts)
+{
+	enum bootstage_id id = BOOTSTAGE_ID_USER + 53;
+	uint start_time, elapsed1, elapsed2;
+	const struct bootstage_record *rec;
+	int index, count;
+
+	count = bootstage_get_rec_count();
+
+	/* Start an accumulator */
+	start_time = bootstage_start(id, "test_accum");
+	ut_assert(start_time > 0);
+
+	/* Check the accumulator record was created */
+	index = count;
+	rec = bootstage_get_rec(index);
+	ut_assertnonnull(rec);
+	ut_asserteq(id, rec->id);
+	ut_asserteq_str("test_accum", rec->name);
+	ut_asserteq(start_time, rec->start_us);
+
+	/* Accumulate the time */
+	udelay(1);
+	elapsed1 = bootstage_accum(id);
+	ut_assert(elapsed1 >= 0);
+
+	/* Check the accumulated time was recorded */
+	ut_asserteq(elapsed1, rec->time_us);
+
+	/* Start and accumulate again  */
+	bootstage_start(id, "test_accum");
+	udelay(1);
+	elapsed2 = bootstage_accum(id);
+	ut_assert(elapsed2 >= 0);
+
+	/* Check the total time accumulated */
+	rec = bootstage_get_rec(index);
+	ut_asserteq(rec->time_us, elapsed1 + elapsed2);
+
+	/* Restore the original count */
+	bootstage_set_rec_count(count);
+
+	return 0;
+}
+COMMON_TEST(test_bootstage_accum, 0);
+
+/* Test bootstage_mark_code() */
+static int test_bootstage_mark_code(struct unit_test_state *uts)
+{
+	const struct bootstage_record *rec;
+	ulong time;
+	int count;
+
+	count = bootstage_get_rec_count();
+
+	/* Mark with file, function, and line number */
+	time = bootstage_mark_code("file.c", __func__, 123);
+	ut_assert(time > 0);
+
+	/* Check the record */
+	rec = bootstage_get_rec(count);
+	ut_assertnonnull(rec);
+	ut_asserteq(time, rec->time_us);
+	ut_asserteq_str("file.c,123: test_bootstage_mark_code", rec->name);
+
+	/* Restore the original count */
+	bootstage_set_rec_count(count);
+
+	return 0;
+}
+COMMON_TEST(test_bootstage_mark_code, 0);
+
+/* Test bootstage_get_rec_count() */
+static int test_bootstage_get_rec_count(struct unit_test_state *uts)
+{
+	const struct bootstage_record *rec;
+	int orig, count;
+
+	/* Get initial count */
+	orig = bootstage_get_rec_count();
+	ut_assert(orig > 0);
+
+	/* Add a new record */
+	bootstage_mark_name(BOOTSTAGE_ID_USER + 52, "test_count");
+
+	/* Verify count increased */
+	count = bootstage_get_rec_count();
+	ut_asserteq(orig + 1, count);
+
+	/* Verify the record was added at the correct index */
+	rec = bootstage_get_rec(orig);
+	ut_assertnonnull(rec);
+	ut_asserteq(BOOTSTAGE_ID_USER + 52, rec->id);
+	ut_asserteq_str("test_count", rec->name);
+
+	/* Restore the original count */
+	bootstage_set_rec_count(orig);
+
+	return 0;
+}
+COMMON_TEST(test_bootstage_get_rec_count, 0);
+
+/* Test bootstage_get_rec() */
+static int test_bootstage_get_rec(struct unit_test_state *uts)
+{
+	const struct bootstage_record *rec;
+	int count;
+
+	/* Get total count */
+	count = bootstage_get_rec_count();
+	ut_assert(count > 0);
+
+	/* Get first record (should be "reset") */
+	rec = bootstage_get_rec(0);
+	ut_assertnonnull(rec);
+	ut_asserteq_str("reset", rec->name);
+
+	/* Test out of bounds access */
+	ut_assertnull(bootstage_get_rec(count));
+	ut_assertnull(bootstage_get_rec(count + 100));
+	ut_assertnull(bootstage_get_rec(-1));
+
+	return 0;
+}
+COMMON_TEST(test_bootstage_get_rec, 0);