helpers.py revision 11482 
1#!/usr/bin/env python
2#
3# Copyright (c) 2016 ARM Limited
4# All rights reserved
5#
6# The license below extends only to copyright in the software and shall
7# not be construed as granting a license to any other intellectual
8# property including but not limited to intellectual property relating
9# to a hardware implementation of the functionality of the software
10# licensed hereunder.  You may use the software subject to the license
11# terms below provided that you ensure that this notice is replicated
12# unmodified and in its entirety in all distributions of the software,
13# modified or unmodified, in source code or in binary form.
14#
15# Redistribution and use in source and binary forms, with or without
16# modification, are permitted provided that the following conditions are
17# met: redistributions of source code must retain the above copyright
18# notice, this list of conditions and the following disclaimer;
19# redistributions in binary form must reproduce the above copyright
20# notice, this list of conditions and the following disclaimer in the
21# documentation and/or other materials provided with the distribution;
22# neither the name of the copyright holders nor the names of its
23# contributors may be used to endorse or promote products derived from
24# this software without specific prior written permission.
25#
26# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
27# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
28# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
29# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
30# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
31# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
32# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
33# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
34# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
35# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
36# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
37#
38# Authors: Andreas Sandberg
39
40import subprocess
41from threading import Timer
42import time
43
44class CallTimeoutException(Exception):
45    """Exception that indicates that a process call timed out"""
46
47    def __init__(self, status, stdout, stderr):
48        self.status = status
49        self.stdout = stdout
50        self.stderr = stderr
51
52class ProcessHelper(subprocess.Popen):
53    """Helper class to run child processes.
54
55    This class wraps a subprocess.Popen class and adds support for
56    using it in a with block. When the process goes out of scope, it's
57    automatically terminated.
58
59    with ProcessHelper(["/bin/ls"], stdout=subprocess.PIPE) as p:
60        return p.call()
61    """
62    def __init__(self, *args, **kwargs):
63        super(ProcessHelper, self).__init__(*args, **kwargs)
64
65    def _terminate_nicely(self, timeout=5):
66        def on_timeout():
67            self.kill()
68
69        if self.returncode is not None:
70            return self.returncode
71
72        timer = Timer(timeout, on_timeout)
73        self.terminate()
74        status = self.wait()
75        timer.cancel()
76
77        return status
78
79    def __enter__(self):
80        return self
81
82    def __exit__(self, exc_type, exc_value, traceback):
83        if self.returncode is None:
84            self._terminate_nicely()
85
86    def call(self, timeout=0):
87        self._timeout = False
88        def on_timeout():
89            self._timeout = True
90            self._terminate_nicely()
91
92        status, stdout, stderr = None, None, None
93        timer = Timer(timeout, on_timeout)
94        if timeout:
95            timer.start()
96
97        stdout, stderr = self.communicate()
98        status = self.wait()
99
100        timer.cancel()
101
102        if self._timeout:
103            self._terminate_nicely()
104            raise CallTimeoutException(self.returncode, stdout, stderr)
105        else:
106            return status, stdout, stderr
107
108if __name__ == "__main__":
109    # Run internal self tests to ensure that the helpers are working
110    # properly. The expected output when running this script is
111    # "SUCCESS!".
112
113    cmd_foo = [ "/bin/echo", "-n", "foo" ]
114    cmd_sleep = [ "/bin/sleep", "10" ]
115
116    # Test that things don't break if the process hasn't been started
117    with ProcessHelper(cmd_foo) as p:
118        pass
119
120    with ProcessHelper(cmd_foo, stdout=subprocess.PIPE) as p:
121        status, stdout, stderr = p.call()
122    assert stdout == "foo"
123    assert status == 0
124
125    try:
126        with ProcessHelper(cmd_sleep) as p:
127            status, stdout, stderr = p.call(timeout=1)
128        assert False, "Timeout not triggered"
129    except CallTimeoutException:
130        pass
131
132    print "SUCCESS!"
133