[U-Boot] test/py: regenerate persistent GPT image if code changes

diff --git a/test/py/tests/test_gpt.py b/test/py/tests/test_gpt.py
index b9b5e5fbb04f..4329b69b7add 100644
--- a/test/py/tests/test_gpt.py
+++ b/test/py/tests/test_gpt.py
@@ -32,23 +32,24 @@  class GptTestDiskImage(object):
         persistent = u_boot_console.config.persistent_data_dir + '/' + filename
         self.path = u_boot_console.config.result_dir  + '/' + filename
 
-        if os.path.exists(persistent):
-            u_boot_console.log.action('Disk image file ' + persistent +
-                ' already exists')
-        else:
-            u_boot_console.log.action('Generating ' + persistent)
-            fd = os.open(persistent, os.O_RDWR | os.O_CREAT)
-            os.ftruncate(fd, 4194304)
-            os.close(fd)
-            cmd = ('sgdisk', '-U', '375a56f7-d6c9-4e81-b5f0-09d41ca89efe',
-                persistent)
-            u_boot_utils.run_and_log(u_boot_console, cmd)
-            cmd = ('sgdisk', '--new=1:2048:2560', '-c 1:part1', persistent)
-            u_boot_utils.run_and_log(u_boot_console, cmd)
-            cmd = ('sgdisk', '--new=2:4096:4608', '-c 2:part2', persistent)
-            u_boot_utils.run_and_log(u_boot_console, cmd)
-            cmd = ('sgdisk', '-l', persistent)
-            u_boot_utils.run_and_log(u_boot_console, cmd)
+        with u_boot_utils.persistent_file_helper(u_boot_console.log, persistent):
+            if os.path.exists(persistent):
+                u_boot_console.log.action('Disk image file ' + persistent +
+                    ' already exists')
+            else:
+                u_boot_console.log.action('Generating ' + persistent)
+                fd = os.open(persistent, os.O_RDWR | os.O_CREAT)
+                os.ftruncate(fd, 4194304)
+                os.close(fd)
+                cmd = ('sgdisk', '-U', '375a56f7-d6c9-4e81-b5f0-09d41ca89efe',
+                    persistent)
+                u_boot_utils.run_and_log(u_boot_console, cmd)
+                cmd = ('sgdisk', '--new=1:2048:2560', '-c 1:part1', persistent)
+                u_boot_utils.run_and_log(u_boot_console, cmd)
+                cmd = ('sgdisk', '--new=2:4096:4608', '-c 2:part2', persistent)
+                u_boot_utils.run_and_log(u_boot_console, cmd)
+                cmd = ('sgdisk', '-l', persistent)
+                u_boot_utils.run_and_log(u_boot_console, cmd)
 
         cmd = ('cp', persistent, self.path)
         u_boot_utils.run_and_log(u_boot_console, cmd)
diff --git a/test/py/u_boot_utils.py b/test/py/u_boot_utils.py
index 2ba4baed07eb..9acb92ddc448 100644
--- a/test/py/u_boot_utils.py
+++ b/test/py/u_boot_utils.py
@@ -5,6 +5,7 @@ 
 # Utility code shared across multiple tests.
 
 import hashlib
+import inspect
 import os
 import os.path
 import pytest
@@ -237,3 +238,76 @@  def find_ram_base(u_boot_console):
             raise Exception('Failed to find RAM bank start in `bdinfo`')
 
     return ram_base
+
+class PersistentFileHelperCtxMgr(object):
+    """A context manager for Python's "with" statement, which ensures that any
+    generated file is deleted (and hence regenerated) if its mtime is older
+    than the mtime of the Python module which generated it, and gets an mtime
+    newer than the mtime of the Python module which generated after it is
+    generated. Objects of this type should be created by factory function
+    persistent_file_helper rather than directly."""
+
+    def __init__(self, log, filename):
+        """Initialize a new object.
+
+        Args:
+            log: The Logfile object to log to.
+            filename: The filename of the generated file.
+
+        Returns:
+            Nothing.
+        """
+
+        self.log = log
+        self.filename = filename
+
+    def __enter__(self):
+        frame = inspect.stack()[1]
+        module = inspect.getmodule(frame[0])
+        self.module_filename = module.__file__
+        self.module_timestamp = os.path.getmtime(self.module_filename)
+
+        if os.path.exists(self.filename):
+            filename_timestamp = os.path.getmtime(self.filename)
+            if filename_timestamp < self.module_timestamp:
+                self.log.action('Removing stale generated file ' +
+                    self.filename)
+                os.unlink(self.filename)
+
+    def __exit__(self, extype, value, traceback):
+        if extype:
+            try:
+                os.path.unlink(self.filename)
+            except:
+                pass
+            return
+        logged = False
+        for i in range(20):
+            filename_timestamp = os.path.getmtime(self.filename)
+            if filename_timestamp > self.module_timestamp:
+                break
+            if not logged:
+                self.log.action(
+                    'Waiting for generated file timestamp to increase')
+                logged = True
+            os.utime(self.filename)
+            time.sleep(0.1)
+
+def persistent_file_helper(u_boot_log, filename):
+    """Manage the timestamps and regeneration of a persistent generated
+    file. This function creates a context manager for Python's "with"
+    statement
+
+    Usage:
+        with persistent_file_helper(u_boot_console.log, filename):
+            code to generate the file, if it's missing.
+
+    Args:
+        u_boot_log: u_boot_console.log.
+        filename: The filename of the generated file.
+
+    Returns:
+        A context manager object.
+    """
+
+    return PersistentFileHelperCtxMgr(u_boot_log, filename)