|  | @@ -1,11 +1,14 @@
 | 
	
		
			
				|  |  |  from __future__ import absolute_import
 | 
	
		
			
				|  |  | +from __future__ import with_statement
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  from mock import patch
 | 
	
		
			
				|  |  | +from contextlib import contexmanager
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  from celery import current_app
 | 
	
		
			
				|  |  | +from celery import result
 | 
	
		
			
				|  |  |  from celery.result import AsyncResult
 | 
	
		
			
				|  |  |  from celery.task import chords
 | 
	
		
			
				|  |  | -from celery.task import TaskSet
 | 
	
		
			
				|  |  | +from celery.task import task, TaskSet
 | 
	
		
			
				|  |  |  from celery.tests.utils import AppCase, Mock
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  passthru = lambda x: x
 | 
	
	
		
			
				|  | @@ -35,47 +38,61 @@ class TSR(chords.TaskSetResult):
 | 
	
		
			
				|  |  |          return self.value
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  | -class test_unlock_chord_task(AppCase):
 | 
	
		
			
				|  |  | +@contextmanager
 | 
	
		
			
				|  |  | +def patch_unlock_retry():
 | 
	
		
			
				|  |  | +    unlock = current_app.tasks["celery.chord_unlock"]
 | 
	
		
			
				|  |  | +    retry = Mock()
 | 
	
		
			
				|  |  | +    prev, unlock.retry = unlock.retry, retry
 | 
	
		
			
				|  |  | +    yield unlock, retry
 | 
	
		
			
				|  |  | +    unlock.retry = prev
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  | -    @patch("celery.task.chords._unlock_chord.retry")
 | 
	
		
			
				|  |  | -    def test_unlock_ready(self, retry):
 | 
	
		
			
				|  |  | -        callback.apply_async = Mock()
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  | -        pts, chords.TaskSetResult = chords.TaskSetResult, TSR
 | 
	
		
			
				|  |  | -        subtask, chords.subtask = chords.subtask, passthru
 | 
	
		
			
				|  |  | -        try:
 | 
	
		
			
				|  |  | -            chords._unlock_chord("setid", callback.subtask(),
 | 
	
		
			
				|  |  | -                    result=map(AsyncResult, [1, 2, 3]))
 | 
	
		
			
				|  |  | -        finally:
 | 
	
		
			
				|  |  | -            chords.subtask = subtask
 | 
	
		
			
				|  |  | -            chords.TaskSetResult = pts
 | 
	
		
			
				|  |  | -        callback.apply_async.assert_called_with(([2, 4, 8, 6], ), {})
 | 
	
		
			
				|  |  | -        # did not retry
 | 
	
		
			
				|  |  | -        self.assertFalse(retry.call_count)
 | 
	
		
			
				|  |  | -
 | 
	
		
			
				|  |  | -    @patch("celery.task.chords.TaskSetResult")
 | 
	
		
			
				|  |  | -    @patch("celery.task.chords._unlock_chord.retry")
 | 
	
		
			
				|  |  | -    def test_when_not_ready(self, retry, TaskSetResult):
 | 
	
		
			
				|  |  | -        callback.apply_async = Mock()
 | 
	
		
			
				|  |  | +class test_unlock_chord_task(AppCase):
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +    @patch("celery.result.TaskSetResult")
 | 
	
		
			
				|  |  | +    def test_unlock_ready(self, TaskSetResult):
 | 
	
		
			
				|  |  | +        tasks = current_app.tasks
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |          class NeverReady(TSR):
 | 
	
		
			
				|  |  |              is_ready = False
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  | -        pts, chords.TaskSetResult = chords.TaskSetResult, NeverReady
 | 
	
		
			
				|  |  | +        @task
 | 
	
		
			
				|  |  | +        def callback(*args, **kwargs):
 | 
	
		
			
				|  |  | +            pass
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +        pts, result.TaskSetResult  = result.TaskSetResult, NeverReady
 | 
	
		
			
				|  |  | +        callback.apply_async = Mock()
 | 
	
		
			
				|  |  |          try:
 | 
	
		
			
				|  |  | -            chords._unlock_chord("setid", callback.subtask, interval=10,
 | 
	
		
			
				|  |  | -                                max_retries=30,
 | 
	
		
			
				|  |  | -                                result=map(AsyncResult, [1, 2, 3]))
 | 
	
		
			
				|  |  | -            self.assertFalse(callback.apply_async.call_count)
 | 
	
		
			
				|  |  | -            # did retry
 | 
	
		
			
				|  |  | -            chords._unlock_chord.retry.assert_called_with(countdown=10,
 | 
	
		
			
				|  |  | -                                                          max_retries=30)
 | 
	
		
			
				|  |  | +            with patch_unlock_retry() as (unlock, retry):
 | 
	
		
			
				|  |  | +                result = Mock(attrs=dict(ready=lambda: True,
 | 
	
		
			
				|  |  | +                                        join=lambda **kw: [2, 4, 8, 6]))
 | 
	
		
			
				|  |  | +                TaskSetResult.restore = lambda setid: result
 | 
	
		
			
				|  |  | +                subtask, chords.subtask = chords.subtask, passthru
 | 
	
		
			
				|  |  | +                try:
 | 
	
		
			
				|  |  | +                    unlock("setid", callback,
 | 
	
		
			
				|  |  | +                           result=map(AsyncResult, [1, 2, 3]))
 | 
	
		
			
				|  |  | +                finally:
 | 
	
		
			
				|  |  | +                    chords.subtask = subtask
 | 
	
		
			
				|  |  | +                callback.apply_async.assert_called_with(([2, 4, 8, 6], ), {})
 | 
	
		
			
				|  |  | +                result.delete.assert_called_with()
 | 
	
		
			
				|  |  | +                # did not retry
 | 
	
		
			
				|  |  | +                self.assertFalse(retry.call_count)
 | 
	
		
			
				|  |  |          finally:
 | 
	
		
			
				|  |  | -            chords.TaskSetResult = pts
 | 
	
		
			
				|  |  | +            result.TaskSetResult = pts
 | 
	
		
			
				|  |  | +
 | 
	
		
			
				|  |  | +    @patch("celery.result.TaskSetResult")
 | 
	
		
			
				|  |  | +    def test_when_not_ready(self, TaskSetResult):
 | 
	
		
			
				|  |  | +        with patch_unlock_retry() as (unlock, retry):
 | 
	
		
			
				|  |  | +            callback = Mock()
 | 
	
		
			
				|  |  | +            result = Mock(attrs=dict(ready=lambda: False))
 | 
	
		
			
				|  |  | +            TaskSetResult.restore = lambda setid: result
 | 
	
		
			
				|  |  | +            unlock("setid", callback, interval=10, max_retries=30,)
 | 
	
		
			
				|  |  | +            self.assertFalse(callback.delay.call_count)
 | 
	
		
			
				|  |  | +            # did retry
 | 
	
		
			
				|  |  | +            unlock.retry.assert_called_with(countdown=10, max_retries=30)
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |      def test_is_in_registry(self):
 | 
	
		
			
				|  |  | -        from celery.registry import tasks
 | 
	
		
			
				|  |  | -        self.assertIn("celery.chord_unlock", tasks)
 | 
	
		
			
				|  |  | +        self.assertIn("celery.chord_unlock", current_app.tasks)
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  
 | 
	
		
			
				|  |  |  class test_chord(AppCase):
 |