--- /dev/null
+# SPDX-License-Identifier: GPL-2.0
+# (C) Copyright 2023, Advanced Micro Devices, Inc.
+
+import pytest
+
+"""
+Test the bootstage command.
+
+It is used for checking the boot progress and timing by printing the bootstage
+report, stashes the data into memory and unstashes the data from memory.
+
+Note: This test relies on boardenv_* containing configuration values to define
+the data size, memory address, and bootstage magic address (defined in
+common/bootstage.c). Without this, bootstage stash and unstash tests will be
+automatically skipped.
+
+For example:
+env__bootstage_cmd_file = {
+ 'addr': 0x200000,
+ 'size': 0x1000,
+ 'bootstage_magic_addr': 0xb00757a3,
+}
+"""
+
+@pytest.mark.buildconfigspec('bootstage')
+@pytest.mark.buildconfigspec('cmd_bootstage')
+def test_bootstage_report(u_boot_console):
+ output = u_boot_console.run_command('bootstage report')
+ assert 'Timer summary in microseconds' in output
+ assert 'Accumulated time:' in output
+ assert 'dm_r' in output
+
+@pytest.mark.buildconfigspec('bootstage')
+@pytest.mark.buildconfigspec('cmd_bootstage')
+@pytest.mark.buildconfigspec('bootstage_stash')
+def test_bootstage_stash(u_boot_console):
+ f = u_boot_console.config.env.get('env__bootstage_cmd_file', None)
+ if not f:
+ pytest.skip('No bootstage environment file is defined')
+
+ addr = f.get('addr')
+ size = f.get('size')
+ bootstage_magic = f.get('bootstage_magic_addr')
+ expected_text = 'dm_r'
+
+ u_boot_console.run_command('bootstage stash %x %x' % (addr, size))
+ output = u_boot_console.run_command('echo $?')
+ assert output.endswith('0')
+
+ output = u_boot_console.run_command('md %x 100' % addr)
+
+ # Check BOOTSTAGE_MAGIC address at 4th byte address
+ assert '0x' + output.split('\n')[0].split()[4] == hex(bootstage_magic)
+
+ # Check expected string in last column of output
+ output_last_col = ''.join([i.split()[-1] for i in output.split('\n')])
+ assert expected_text in output_last_col
+ return addr, size
+
+@pytest.mark.buildconfigspec('bootstage')
+@pytest.mark.buildconfigspec('cmd_bootstage')
+@pytest.mark.buildconfigspec('bootstage_stash')
+def test_bootstage_unstash(u_boot_console):
+ addr, size = test_bootstage_stash(u_boot_console)
+ u_boot_console.run_command('bootstage unstash %x %x' % (addr, size))
+ output = u_boot_console.run_command('echo $?')
+ assert output.endswith('0')