1 files changed, 132 insertions, 0 deletions
diff --git a/tests/testing/helpers.py b/tests/testing/helpers.py
new file mode 100755
index 000000000..dcc48904c
--- /dev/null
+++ b/tests/testing/helpers.py
@@ -0,0 +1,132 @@
+#!/usr/bin/env python
+#
+# Copyright (c) 2016 ARM Limited
+# All rights reserved
+#
+# The license below extends only to copyright in the software and shall
+# not be construed as granting a license to any other intellectual
+# property including but not limited to intellectual property relating
+# to a hardware implementation of the functionality of the software
+# licensed hereunder.  You may use the software subject to the license
+# terms below provided that you ensure that this notice is replicated
+# unmodified and in its entirety in all distributions of the software,
+# modified or unmodified, in source code or in binary form.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are
+# met: redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer;
+# redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution;
+# neither the name of the copyright holders nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# Authors: Andreas Sandberg
+
+import subprocess
+from threading import Timer
+import time
+
+class CallTimeoutException(Exception):
+    """Exception that indicates that a process call timed out"""
+
+    def __init__(self, status, stdout, stderr):
+        self.status = status
+        self.stdout = stdout
+        self.stderr = stderr
+
+class ProcessHelper(subprocess.Popen):
+    """Helper class to run child processes.
+
+    This class wraps a subprocess.Popen class and adds support for
+    using it in a with block. When the process goes out of scope, it's
+    automatically terminated.
+
+    with ProcessHelper(["/bin/ls"], stdout=subprocess.PIPE) as p:
+        return p.call()
+    """
+    def __init__(self, *args, **kwargs):
+        super(ProcessHelper, self).__init__(*args, **kwargs)
+
+    def _terminate_nicely(self, timeout=5):
+        def on_timeout():
+            self.kill()
+
+        if self.returncode is not None:
+            return self.returncode
+
+        timer = Timer(timeout, on_timeout)
+        self.terminate()
+        status = self.wait()
+        timer.cancel()
+
+        return status
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.returncode is None:
+            self._terminate_nicely()
+
+    def call(self, timeout=0):
+        self._timeout = False
+        def on_timeout():
+            self._timeout = True
+            self._terminate_nicely()
+
+        status, stdout, stderr = None, None, None
+        timer = Timer(timeout, on_timeout)
+        if timeout:
+            timer.start()
+
+        stdout, stderr = self.communicate()
+        status = self.wait()
+
+        timer.cancel()
+
+        if self._timeout:
+            self._terminate_nicely()
+            raise CallTimeoutException(self.returncode, stdout, stderr)
+        else:
+            return status, stdout, stderr
+
+if __name__ == "__main__":
+    # Run internal self tests to ensure that the helpers are working
+    # properly. The expected output when running this script is
+    # "SUCCESS!".
+
+    cmd_foo = [ "/bin/echo", "-n", "foo" ]
+    cmd_sleep = [ "/bin/sleep", "10" ]
+
+    # Test that things don't break if the process hasn't been started
+    with ProcessHelper(cmd_foo) as p:
+        pass
+
+    with ProcessHelper(cmd_foo, stdout=subprocess.PIPE) as p:
+        status, stdout, stderr = p.call()
+    assert stdout == "foo"
+    assert status == 0
+
+    try:
+        with ProcessHelper(cmd_sleep) as p:
+            status, stdout, stderr = p.call(timeout=1)
+        assert False, "Timeout not triggered"
+    except CallTimeoutException:
+        pass
+
+    print "SUCCESS!"