2 files changed, 32 insertions, 0 deletions
diff --git a/test/py/u_boot_console_base.py b/test/py/u_boot_console_base.py
index f7434363fb..815fa64d5f 100644
--- a/test/py/u_boot_console_base.py
+++ b/test/py/u_boot_console_base.py
@@ -56,6 +56,22 @@ class ConsoleDisableCheck(object):
         self.console.disable_check_count[self.check_type] -= 1
         self.console.eval_bad_patterns()
 
+class ConsoleSetupTimeout(object):
+    """Context manager (for Python's with statement) that temporarily sets up
+    timeout for specific command. This is useful when execution time is greater
+    then default 30s."""
+
+    def __init__(self, console, timeout):
+        self.p = console.p
+        self.orig_timeout = self.p.timeout
+        self.p.timeout = timeout
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, extype, value, traceback):
+        self.p.timeout = self.orig_timeout
+
 class ConsoleBase(object):
     """The interface through which test functions interact with the U-Boot
     console. This primarily involves executing shell commands, capturing their
@@ -391,3 +407,18 @@ class ConsoleBase(object):
         """
 
         return ConsoleDisableCheck(self, check_type)
+
+    def temporary_timeout(self, timeout):
+        """Temporarily set up different timeout for commands.
+
+        Create a new context manager (for use with the "with" statement) which
+        temporarily change timeout.
+
+        Args:
+            timeout: Time in milliseconds.
+
+        Returns:
+            A context manager object.
+        """
+
+        return ConsoleSetupTimeout(self, timeout)
diff --git a/test/py/u_boot_utils.py b/test/py/u_boot_utils.py
index 9d243e047b..6a6b2ec0e6 100644
--- a/test/py/u_boot_utils.py
+++ b/test/py/u_boot_utils.py
@@ -10,6 +10,7 @@ import os.path
 import pytest
 import sys
 import time
+import pytest
 
 def md5sum_data(data):
     """Calculate the MD5 hash of some data.