diff mbox series

[Concept,19/19] test: Save and restore bootstage record count in FIT tests

Message ID 20260314231618.338113-20-sjg@u-boot.org
State New
Headers
Series test: Fix pytest inter-test side effects |

Commit Message

Simon Glass March 14, 2026, 11:16 p.m. UTC 
  From: Simon Glass <sjg@chromium.org>

Add a preserve_bootstage() context-manager helper in the pytest
utils module, and use it in TestFitImage which runs bootm and is the
main consumer of bootstage records.

Also save and restore the bootstage record count in test_pre_run()
and test_post_run() for ut tests which run multiple tests within a
single 'ut' command invocation.

Signed-off-by: Simon Glass <sjg@chromium.org>
---

 include/test/test.h       |  2 ++
 test/py/tests/test_fit.py |  6 ++++++
 test/py/utils.py          | 18 ++++++++++++++++++
 test/test-main.c          |  5 +++++
 4 files changed, 31 insertions(+)
diff mbox series

Patch

diff --git a/include/test/test.h b/include/test/test.h
index c3b251e2cd4..b81cab4d7a4 100644
--- a/include/test/test.h
+++ b/include/test/test.h
@@ -92,6 +92,7 @@  struct ut_arg {
  * @force_run: true to run tests marked with the UTF_MANUAL flag
  * @workers: Number of parallel workers, 0 if not sharding tests
  * @worker_id: ID of this worker (0 to workers-1)
+ * @old_bootstage_count: bootstage record count saved before each test
  * @old_bloblist: stores the old gd->bloblist pointer
  * @soft_fail: continue execution of the test even after it fails
  * @expect_str: Temporary string used to hold expected string value
@@ -128,6 +129,7 @@  struct unit_test_state {
 	bool force_run;
 	int workers;
 	int worker_id;
+	uint old_bootstage_count;
 	void *old_bloblist;
 	bool soft_fail;
 	char expect_str[1024];
diff --git a/test/py/tests/test_fit.py b/test/py/tests/test_fit.py
index ed18ff68825..6fb76196e16 100755
--- a/test/py/tests/test_fit.py
+++ b/test/py/tests/test_fit.py
@@ -140,6 +140,12 @@  class TestFitImage:
       - run code coverage to make sure we are testing all the code
     """
 
+    @pytest.fixture(autouse=True)
+    def save_bootstage(self, ubman):
+        """Save and restore bootstage around each test."""
+        with utils.preserve_bootstage(ubman):
+            yield
+
     def make_fname(self, ubman, leaf):
         """Make a temporary filename
 
diff --git a/test/py/utils.py b/test/py/utils.py
index 083a3feddc2..9097f158496 100644
--- a/test/py/utils.py
+++ b/test/py/utils.py
@@ -14,8 +14,26 @@  import signal
 import sys
 import time
 import re
+from contextlib import contextmanager
 import pytest
 
+@contextmanager
+def preserve_bootstage(ubman):
+    """Context manager to save and restore bootstage record count.
+
+    Some commands (e.g. bootm) add bootstage records with unique IDs. These
+    accumulate across tests in a pytest session and can fill the bootstage
+    table. Use this around tests that trigger such commands.
+
+    Args:
+        ubman (ConsoleBase): U-Boot console connection
+    """
+    ubman.run_command('bootstage save')
+    try:
+        yield
+    finally:
+        ubman.run_command('bootstage restore')
+
 def md5sum_data(data):
     """Calculate the MD5 hash of some data.
 
diff --git a/test/test-main.c b/test/test-main.c
index 5db35b59760..77223cfbcb7 100644
--- a/test/test-main.c
+++ b/test/test-main.c
@@ -7,6 +7,7 @@ 
 #define LOG_CATEGORY	LOGC_TEST
 
 #include <blk.h>
+#include <bootstage.h>
 #include <console.h>
 #include <cyclic.h>
 #include <dm.h>
@@ -560,6 +561,8 @@  static int test_pre_run(struct unit_test_state *uts, struct unit_test *test)
 		gd_set_bloblist(NULL);
 	}
 
+	uts->old_bootstage_count = bootstage_get_rec_count();
+
 	if (!(test->flags & UTF_NO_SILENT))
 		ut_silence_console(uts);
 
@@ -589,6 +592,8 @@  static int test_post_run(struct unit_test_state *uts, struct unit_test *test)
 		log_debug("restore bloblist %p\n", gd_bloblist());
 	}
 
+	bootstage_set_rec_count(uts->old_bootstage_count);
+
 	blkcache_free();
 
 	return 0;