fix apply

kuyruk/task.py

@@ -62,14 +62,13 @@ def __call__(self, *args, **kwargs):
         """
         self.send_signal(events.task_presend, args, kwargs, reverse=True)
 
-        task_result = TaskResult(self)
-
         host = kwargs.pop('kuyruk_host', None)
         local = kwargs.pop('kuyruk_local', None)
 
         if self.eager or self.kuyruk.config.EAGER:
-            task_result.result = self.apply(*args, **kwargs)
+            task_result = self.apply(*args, **kwargs)
         else:
+            task_result = TaskResult(self)
             task_result.id = self.send_to_queue(args, kwargs,
                                                 host=host, local=local)
 
@@ -169,6 +168,8 @@ def apply(self, *args, **kwargs):
         def send_signal(signal, reverse=False, **extra):
             self.send_signal(signal, args, kwargs, reverse, **extra)
 
+        result = TaskResult(self)
+
         limit = (self.max_run_time or
                  self.kuyruk.config.MAX_TASK_RUN_TIME or 0)
 
@@ -183,9 +184,12 @@ def send_signal(signal, reverse=False, **extra):
             raise
         else:
             send_signal(events.task_success, return_value=return_value)
+            result.result = return_value
         finally:
             send_signal(events.task_postrun)
 
+        return result
+
     @property
     def name(self):
         """Location for the wrapped function.

kuyruk/test/tasks.py

@@ -12,6 +12,7 @@
 @kuyruk.task
 def print_task(message):
     print message
+    must_be_called()
 
 
 @kuyruk.task(queue='another_queue')
@@ -128,6 +129,10 @@ def jump(cat):
 
 
 def must_be_called():
+    """
+    This function is patched in tests to see the caller is doing it's job.
+
+    """
     print 'Yes, it is called.'
 
 

kuyruk/test/test_kuyruk.py

@@ -78,6 +78,13 @@ def test_eager(self, mock_func):
         assert isinstance(result, TaskResult)
         mock_func.assert_called_once_with()
 
+    @patch('kuyruk.test.tasks.must_be_called')
+    def test_apply(self, mock_func):
+        """Test Task.apply()"""
+        result = tasks.print_task.apply("hello")
+        assert isinstance(result, TaskResult)
+        mock_func.assert_called_once_with()
+
     def test_reject(self):
         """Rejected tasks must be requeued again"""
         tasks.rejecting_task()