[tor-commits] [ooni-probe/master] Add support for specifying a timeout after which to close the process

[art at torproject.org]

commit 0d170c31418f27aea936dcd9148e18ec10994c9f
Author: Arturo Filastò <art at fuffa.org>
Date:   Thu Mar 26 14:08:34 2015 +0100

    Add support for specifying a timeout after which to close the process
    
    * Better exit reason handling
---
 data/nettests/examples/example_process.py |    1 +
 ooni/templates/process.py                 |   61 +++++++++++++++++++----------
 2 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/data/nettests/examples/example_process.py b/data/nettests/examples/example_process.py
index 32e52e1..a07f5f1 100644
--- a/data/nettests/examples/example_process.py
+++ b/data/nettests/examples/example_process.py
@@ -7,3 +7,4 @@ class TestProcessExample(process.ProcessTest):
     @defer.inlineCallbacks
     def test_http_and_dns(self):
         yield self.run(["echo", "Hello world!"])
+        yield self.run(["sleep", "10"])
diff --git a/ooni/templates/process.py b/ooni/templates/process.py
index 736dbcd..c73ea29 100644
--- a/ooni/templates/process.py
+++ b/ooni/templates/process.py
@@ -1,26 +1,40 @@
 from twisted.internet import protocol, defer, reactor
 
 from ooni.nettest import NetTestCase
-from ooni.errors import failureToString
 from ooni.utils import log
 
 
 class ProcessDirector(protocol.ProcessProtocol):
-    def __init__(self, d, finished=None):
+    def __init__(self, d, finished=None, timeout=None, stdin=None):
         self.d = d
         self.stderr = ""
         self.stdout = ""
         self.finished = finished
-
-    def data(self):
-        return
-
-    def finish(self, reason=None):
-        exit_reason = failureToString(reason)
+        self.timeout = timeout
+        self.stdin = stdin
+
+        self.timer = None
+        self.exit_reason = None
+
+    def close(self, reason=None):
+        self.reason = reason
+        self.transport.loseConnection()
+
+    def resetTimer(self):
+        if self.timeout is not None:
+            if self.timer is not None:
+                self.timer.cancel()
+            self.timer = reactor.callLater(self.timeout,
+                                           self.close,
+                                           "timeout_reached")
+
+    def finish(self, exit_reason=None):
+        if not self.exit_reason:
+            self.exit_reason = exit_reason
         data = {
             "stderr": self.stderr,
             "stdout": self.stdout,
-            "exit_reason": exit_reason
+            "exit_reason": self.exit_reason
         }
         self.d.callback(data)
 
@@ -30,14 +44,17 @@ class ProcessDirector(protocol.ProcessProtocol):
         return self.finished(self.stdout, self.stderr)
 
     def connectionMade(self):
-        self.transport.write("")
-        self.transport.closeStdin()
+        self.resetTimer()
+        if self.stdin is not None:
+            self.transport.write(self.stin)
+            self.transport.closeStdin()
 
     def outReceived(self, data):
+        self.resetTimer()
         log.debug("STDOUT: %s" % data)
         self.stdout += data
         if self.shouldClose():
-            self.transport.loseConnection()
+            self.close("condition_met")
 
     def errReceived(self, data):
         log.debug("STDERR: %s" % data)
@@ -58,7 +75,7 @@ class ProcessDirector(protocol.ProcessProtocol):
 
     def processEnded(self, reason):
         log.debug("Ended %s" % reason)
-        self.finish(reason)
+        self.finish("process_done")
 
 
 class ProcessTest(NetTestCase):
@@ -67,19 +84,23 @@ class ProcessTest(NetTestCase):
 
     requiresRoot = False
     timeout = 5
-    address = None
-    port = None
 
     def _setUp(self):
         super(ProcessTest, self)._setUp()
 
-    def processEnded(self, result):
-        self.report.update(result)
+    def processEnded(self, result, command):
+        log.debug("Finished %s: %s" % (command, result))
+        key = ' '.join(command)
+        self.report[key] = {
+            'stdout': result['stdout'],
+            'stderr': result['stderr'],
+            'exit_reason': result['exit_reason']
+        }
         return result
 
-    def run(self, command):
+    def run(self, command, finished=None):
         d = defer.Deferred()
-        d.addCallback(self.processEnded)
-        processDirector = ProcessDirector(d)
+        d.addCallback(self.processEnded, command)
+        processDirector = ProcessDirector(d, finished, self.timeout)
         reactor.spawnProcess(processDirector, command[0], command)
         return d