[tor-commits] [ooni-probe/master] Implement and unittest code for tasks that timeout

[isis at torproject.org]

commit ba8fb7ce88d179bf7f3ae09017b6ec3b5eca3f96
Author: Arturo Filastò <art at fuffa.org>
Date:   Sun Jan 13 00:09:17 2013 +0100

    Implement and unittest code for tasks that timeout
---
 ooni/tasks.py          |   19 +++++++-----
 tests/test_managers.py |   78 ++++++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 80 insertions(+), 17 deletions(-)

diff --git a/ooni/tasks.py b/ooni/tasks.py
index 703aec3..3230732 100644
--- a/ooni/tasks.py
+++ b/ooni/tasks.py
@@ -21,10 +21,10 @@ class BaseTask(object):
         return result
 
     def start(self):
-        d = self.run()
-        d.addErrback(self._failed)
-        d.addCallback(self._succeeded)
-        return d
+        self.running = self.run()
+        self.running.addErrback(self._failed)
+        self.running.addCallback(self._succeeded)
+        return self.running
 
     def succeeded(self, result):
         """
@@ -50,13 +50,17 @@ class TaskTimedOut(Exception):
 
 class TaskWithTimeout(BaseTask):
     timeout = 5
+    # So that we can test the callLater calls
+    clock = reactor
 
     def _timedOut(self):
         """Internal method for handling timeout failure"""
         self.timedOut()
+        self.running.errback(TaskTimedOut)
 
     def _cancelTimer(self):
-        if self._timer:
+        #import pdb; pdb.set_trace()
+        if self._timer.active():
             self._timer.cancel()
 
     def _succeeded(self, result):
@@ -68,9 +72,8 @@ class TaskWithTimeout(BaseTask):
         return BaseTask._failed(self, failure)
 
     def start(self):
-        self._timer = reactor.callLater(self.timeout, self._timedOut)
-        d = BaseTask.start(self)
-        return d
+        self._timer = self.clock.callLater(self.timeout, self._timedOut)
+        return BaseTask.start(self)
 
     def timedOut(self):
         """
diff --git a/tests/test_managers.py b/tests/test_managers.py
index a10933e..d4ee2b5 100644
--- a/tests/test_managers.py
+++ b/tests/test_managers.py
@@ -1,6 +1,6 @@
 from twisted.trial import unittest
 from twisted.python import failure
-from twisted.internet import defer
+from twisted.internet import defer, task
 
 from ooni.tasks import BaseTask, TaskWithTimeout
 from ooni.managers import TaskManager, MeasurementManager
@@ -26,6 +26,17 @@ class MockSuccessTaskWithTimeout(TaskWithTimeout):
     def run(self):
         return defer.succeed(42)
 
+class MockFailTaskThatTimesOut(TaskWithTimeout):
+    def run(self):
+        return defer.Deferred()
+
+class MockTimeoutOnceTask(TaskWithTimeout):
+    def run(self):
+        if self.failures >= 1:
+            return defer.succeed(42)
+        else:
+            return defer.Deferred()
+
 class MockFailTaskWithTimeout(TaskWithTimeout):
     def run(self):
         return defer.fail(mockFailure)
@@ -48,7 +59,9 @@ class TestTaskManager(unittest.TestCase):
 
         self.measurementManager.start()
 
-    def schedule_failing_tasks(self, task_type, number=1):
+        self.clock = task.Clock()
+
+    def schedule_successful_tasks(self, task_type, number=1):
         all_done = []
         for x in range(number):
             mock_task = task_type()
@@ -58,15 +71,13 @@ class TestTaskManager(unittest.TestCase):
         d = defer.DeferredList(all_done)
         @d.addCallback
         def done(res):
-            # 10*2 because 2 is the number of retries
-            self.assertEqual(len(self.measurementManager.failures), number*3)
-            for task_result, task_instance in self.measurementManager.failures:
-                self.assertEqual(task_result, mockFailure)
+            for task_result, task_instance in self.measurementManager.successes:
+                self.assertEqual(task_result, 42)
                 self.assertIsInstance(task_instance, task_type)
 
         return d
 
-    def schedule_successful_tasks(self, task_type, number=1):
+    def schedule_failing_tasks(self, task_type, number=1):
         all_done = []
         for x in range(number):
             mock_task = task_type()
@@ -76,8 +87,10 @@ class TestTaskManager(unittest.TestCase):
         d = defer.DeferredList(all_done)
         @d.addCallback
         def done(res):
-            for task_result, task_instance in self.measurementManager.successes:
-                self.assertEqual(task_result, 42)
+            # 10*2 because 2 is the number of retries
+            self.assertEqual(len(self.measurementManager.failures), number*3)
+            for task_result, task_instance in self.measurementManager.failures:
+                self.assertEqual(task_result, mockFailure)
                 self.assertIsInstance(task_instance, task_type)
 
         return d
@@ -88,6 +101,53 @@ class TestTaskManager(unittest.TestCase):
     def test_schedule_successful_one_task_with_timeout(self):
         return self.schedule_successful_tasks(MockSuccessTaskWithTimeout)
 
+    def test_schedule_failing_tasks_that_timesout(self):
+        self.measurementManager.retries = 0
+
+        task_type = MockFailTaskThatTimesOut
+        task_timeout = 5
+
+        mock_task = task_type()
+        mock_task.timeout = task_timeout
+        mock_task.clock = self.clock
+
+        self.measurementManager.schedule(mock_task)
+
+        self.clock.advance(task_timeout)
+
+        @mock_task.done.addBoth
+        def done(res):
+            self.assertEqual(len(self.measurementManager.failures), 1)
+            for task_result, task_instance in self.measurementManager.failures:
+                self.assertIsInstance(task_instance, task_type)
+
+        return mock_task.done
+
+    def test_schedule_time_out_once(self):
+        task_type = MockTimeoutOnceTask
+        task_timeout = 5
+
+        mock_task = task_type()
+        mock_task.timeout = task_timeout
+        mock_task.clock = self.clock
+
+        self.measurementManager.schedule(mock_task)
+
+        self.clock.advance(task_timeout)
+
+        @mock_task.done.addBoth
+        def done(res):
+            self.assertEqual(len(self.measurementManager.failures), 1)
+            for task_result, task_instance in self.measurementManager.failures:
+                self.assertIsInstance(task_instance, task_type)
+
+            for task_result, task_instance in self.measurementManager.successes:
+                self.assertEqual(task_result, 42)
+                self.assertIsInstance(task_instance, task_type)
+
+        return mock_task.done
+
+
     def test_schedule_failing_one_task(self):
         return self.schedule_failing_tasks(MockFailTask)