test/py/tests/test_bootstage.py: Combine stash/unstash tests

When running the bootstage tests currently we get a warning like:
tests/test_bootstage.py::test_bootstage_stash
...PytestReturnNotNoneWarning: Expected None, but tests/test_bootstage.py::test_bootstage_stash returned (37748736, 4096), which will be an error in a future version of pytest.  Did you mean to use `assert` in stead of `return`?

This is because the unstash test will run the stash test and fetch the
addr / size from that. Rework the test to be stash and unstash and then
run the unstash command at the end of the current stash test.

Acked-by: Michal Simek <michal.simek@amd.com>
Reviewed-by: Love Kumar <love.kumar@amd.com>
Reviewed-by: Simon Glass <sjg@chromium.org>
Signed-off-by: Tom Rini <trini@konsulko.com>
diff --git a/test/py/tests/test_bootstage.py b/test/py/tests/test_bootstage.py
index a9eb9f0..bd71a1a 100644
--- a/test/py/tests/test_bootstage.py
+++ b/test/py/tests/test_bootstage.py
@@ -33,7 +33,7 @@
 @pytest.mark.buildconfigspec('bootstage')
 @pytest.mark.buildconfigspec('cmd_bootstage')
 @pytest.mark.buildconfigspec('bootstage_stash')
-def test_bootstage_stash(u_boot_console):
+def test_bootstage_stash_and_unstash(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')
@@ -55,13 +55,8 @@
     # 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)
+    # Check that unstash works as expected
     u_boot_console.run_command('bootstage unstash %x %x' % (addr, size))
     output = u_boot_console.run_command('echo $?')
     assert output.endswith('0')