100% coverage for celery.concurrency.gevent

celery/concurrency/gevent.py

@@ -15,7 +15,7 @@ if not os.environ.get('GEVENT_NOPATCH') and not PATCHED[0]:
     PATCHED[0] += 1
     from gevent import monkey, version_info
     monkey.patch_all()
-    if version_info[0] == 0:
+    if version_info[0] == 0:  # pragma: no cover
         # Signals are not working along gevent in version prior 1.0
         # and they are not monkey patch by monkey.patch_all()
         from gevent import signal as _gevent_signal
@@ -24,7 +24,7 @@ if not os.environ.get('GEVENT_NOPATCH') and not PATCHED[0]:
 
 try:
     from gevent import Timeout
-except ImportError:
+except ImportError:  # pragma: no cover
     Timeout = None  # noqa
 
 from time import time
@@ -36,7 +36,8 @@ from .base import apply_target, BasePool
 
 def apply_timeout(target, args=(), kwargs={}, callback=None,
                   accept_callback=None, pid=None, timeout=None,
-                  timeout_callback=None, **rest):
+                  timeout_callback=None, Timeout=Timeout,
+                  apply_target=apply_target, **rest):
     try:
         with Timeout(timeout):
             return apply_target(target, args, kwargs, callback,
@@ -51,9 +52,7 @@ class Schedule(timer2.Schedule):
         from gevent.greenlet import Greenlet, GreenletExit
 
         class _Greenlet(Greenlet):
-
-            def cancel(self):
-                self.kill()
+            cancel = Greenlet.kill
 
         self._Greenlet = _Greenlet
         self._GreenletExit = GreenletExit

celery/tests/concurrency/test_gevent.py

@@ -10,9 +10,13 @@ from celery.concurrency.gevent import (
     Schedule,
     Timer,
     TaskPool,
+    apply_timeout,
+)
+
+from celery.tests.case import (
+    Case, mock_module, patch_many, skip_if_pypy,
 )
 
-from celery.tests.case import Case, mock_module, patch_many, skip_if_pypy
 gevent_modules = (
     'gevent',
     'gevent.monkey',
@@ -80,6 +84,9 @@ class test_Schedule(Case):
                 g.kill.side_effect = KeyError()
                 x.clear()
 
+                g = x._Greenlet()
+                g.cancel()
+
 
 class test_TasKPool(Case):
 
@@ -118,3 +125,37 @@ class test_Timer(Case):
             x.start()
             x.stop()
             x.schedule.clear.assert_called_with()
+
+
+class test_apply_timeout(Case):
+
+    def test_apply_timeout(self):
+
+            class Timeout(Exception):
+                value = None
+
+                def __init__(self, value):
+                    self.__class__.value = value
+
+                def __enter__(self):
+                    return self
+
+                def __exit__(self, *exc_info):
+                    pass
+            timeout_callback = Mock(name='timeout_callback')
+            apply_target = Mock(name='apply_target')
+            apply_timeout(
+                Mock(), timeout=10, callback=Mock(name='callback'),
+                timeout_callback=timeout_callback,
+                apply_target=apply_target, Timeout=Timeout,
+            )
+            self.assertEqual(Timeout.value, 10)
+            self.assertTrue(apply_target.called)
+
+            apply_target.side_effect = Timeout(10)
+            apply_timeout(
+                Mock(), timeout=10, callback=Mock(),
+                timeout_callback=timeout_callback,
+                apply_target=apply_target, Timeout=Timeout,
+            )
+            timeout_callback.assert_called_with(False, 10)