94% total coverage

Author: Ask Solem

celery/backends/cache.py

@@ -6,20 +6,40 @@
 from celery.backends.base import KeyValueStoreBackend
 from celery.exceptions import ImproperlyConfigured
 from celery.utils import timeutils
+from celery.datastructures import LocalCache
 
-try:
-    import pylibmc as memcache
-except ImportError:
+
+
+def get_best_memcache(*args, **kwargs):
     try:
-        import memcache
+        import pylibmc as memcache
     except ImportError:
-        raise ImproperlyConfigured("Memcached backend requires either "
-                                   "the 'memcache' or 'pylibmc' library")
+        try:
+            import memcache
+        except ImportError:
+            raise ImproperlyConfigured("Memcached backend requires either "
+                                       "the 'memcache' or 'pylibmc' library")
+    return memcache.Client(*args, **kwargs)
 
 
-class CacheBackend(KeyValueStoreBackend):
-    Client = memcache.Client
+class DummyClient(object):
+
+    def __init__(self, *args, **kwargs):
+        self.cache = LocalCache(5000)
+
+    def get(self, key, *args, **kwargs):
+        return self.cache.get(key)
 
+    def set(self, key, value, *args, **kwargs):
+        self.cache[key] = value
+
+
+backends = {"memcache": get_best_memcache,
+            "memcached": get_best_memcache,
+            "pylibmc": get_best_memcache,
+            "memory": DummyClient}
+
+class CacheBackend(KeyValueStoreBackend):
     _client = None
 
     def __init__(self, expires=conf.TASK_RESULT_EXPIRES,
@@ -31,6 +51,14 @@ def __init__(self, expires=conf.TASK_RESULT_EXPIRES,
         self.options = dict(conf.CACHE_BACKEND_OPTIONS, **options)
         self.backend, _, servers = partition(backend, "://")
         self.servers = servers.split(";")
+        try:
+            self.Client = backends[self.backend]
+        except KeyError:
+            raise ImproperlyConfigured(
+                    "Unknown cache backend: %s. Please use one of the "
+                    "following backends: %s" % (self.backend,
+                                                ", ".join(backends.keys())))
+
 
     def get(self, key):
         return self.client.get(key)

celery/backends/database.py

@@ -44,18 +44,6 @@ def _store_result(self, task_id, result, status, traceback=None):
             session.close()
         return result
 
-    def _save_taskset(self, taskset_id, result):
-        """Store the result of an executed taskset."""
-        taskset = TaskSet(taskset_id, result)
-        session = self.ResultSession()
-        try:
-            session.add(taskset)
-            session.flush()
-            session.commit()
-        finally:
-            session.close()
-        return result
-
     def _get_task_meta_for(self, task_id):
         """Get task metadata for a task by id."""
         session = self.ResultSession()
@@ -66,17 +54,28 @@ def _get_task_meta_for(self, task_id):
                 session.add(task)
                 session.flush()
                 session.commit()
-            if task:
-                return task.to_dict()
+            return task.to_dict()
+        finally:
+            session.close()
+
+    def _save_taskset(self, taskset_id, result):
+        """Store the result of an executed taskset."""
+        session = self.ResultSession()
+        try:
+            taskset = TaskSet(taskset_id, result)
+            session.add(taskset)
+            session.flush()
+            session.commit()
+            return result
         finally:
             session.close()
 
     def _restore_taskset(self, taskset_id):
         """Get taskset metadata for a taskset by id."""
         session = self.ResultSession()
         try:
-            qs = session.query(TaskSet)
-            taskset = qs.filter(TaskSet.taskset_id == taskset_id).first()
+            taskset = session.query(TaskSet).filter(
+                    TaskSet.taskset_id == taskset_id).first()
             if taskset:
                 return taskset.to_dict()
         finally:

celery/db/models.py

@@ -26,22 +26,14 @@ class Task(ResultModelBase):
     def __init__(self, task_id):
         self.task_id = task_id
 
-    def __str__(self):
-        return "<Task(%s, %s, %s, %s)>" % (self.task_id,
-                                           self.result,
-                                           self.status,
-                                           self.traceback)
-
     def to_dict(self):
         return {"task_id": self.task_id,
                 "status": self.status,
                 "result": self.result,
-                "date_done": self.date_done,
                 "traceback": self.traceback}
 
-    def __unicode__(self):
-        return u"<Task: %s successful: %s>" % (self.task_id, self.status)
-
+    def __repr__(self):
+        return "<Task %s state: %s>" % (self.task_id, self.status)
 
 class TaskSet(ResultModelBase):
     """TaskSet result"""
@@ -55,16 +47,13 @@ class TaskSet(ResultModelBase):
     date_done = sa.Column(sa.DateTime, default=datetime.now,
                        nullable=True)
 
-    def __init__(self, task_id):
-        self.task_id = task_id
-
-    def __str__(self):
-        return "<TaskSet(%s, %s)>" % (self.task_id, self.result)
+    def __init__(self, taskset_id, result):
+        self.taskset_id = taskset_id
+        self.result = result
 
     def to_dict(self):
         return {"taskset_id": self.taskset_id,
-                "result": self.result,
-                "date_done": self.date_done}
+                "result": self.result}
 
-    def __unicode__(self):
+    def __repr__(self):
         return u"<TaskSet: %s>" % (self.taskset_id, )

celery/events/state.py

@@ -4,29 +4,19 @@
 from carrot.utils import partition
 
 from celery import states
-from celery.datastructures import LocalCache
+from celery.datastructures import AttributeDict, LocalCache
 from celery.utils import kwdict
 
 HEARTBEAT_EXPIRE = 150 # 2 minutes, 30 seconds
 
 
-class Element(dict):
+class Element(AttributeDict):
     """Base class for types."""
     visited = False
 
     def __init__(self, **fields):
         dict.__init__(self, fields)
 
-    def __getattr__(self, key):
-        try:
-            return self[key]
-        except KeyError:
-            raise AttributeError("'%s' object has no attribute '%s'" % (
-                    self.__class__.__name__, key))
-
-    def __setattr__(self, key, value):
-        self[key] = value
-
 
 class Worker(Element):
     """Worker State."""

celery/tests/test_backends/test_cache.py

@@ -0,0 +1,131 @@
+import sys
+import types
+import unittest2 as unittest
+
+from celery import states
+from celery.backends.cache import CacheBackend, get_best_memcache, DummyClient
+from celery.exceptions import ImproperlyConfigured
+from celery.utils import gen_unique_id
+
+from celery.tests.utils import execute_context, mask_modules
+
+
+class SomeClass(object):
+
+    def __init__(self, data):
+        self.data = data
+
+
+class test_CacheBackend(unittest.TestCase):
+
+    def test_mark_as_done(self):
+        tb = CacheBackend(backend="memory://")
+
+        tid = gen_unique_id()
+
+        self.assertEqual(tb.get_status(tid), states.PENDING)
+        self.assertIsNone(tb.get_result(tid))
+
+        tb.mark_as_done(tid, 42)
+        self.assertEqual(tb.get_status(tid), states.SUCCESS)
+        self.assertEqual(tb.get_result(tid), 42)
+
+    def test_is_pickled(self):
+        tb = CacheBackend(backend="memory://")
+
+        tid2 = gen_unique_id()
+        result = {"foo": "baz", "bar": SomeClass(12345)}
+        tb.mark_as_done(tid2, result)
+        # is serialized properly.
+        rindb = tb.get_result(tid2)
+        self.assertEqual(rindb.get("foo"), "baz")
+        self.assertEqual(rindb.get("bar").data, 12345)
+
+    def test_mark_as_failure(self):
+        tb = CacheBackend(backend="memory://")
+
+        tid3 = gen_unique_id()
+        try:
+            raise KeyError("foo")
+        except KeyError, exception:
+            pass
+        tb.mark_as_failure(tid3, exception)
+        self.assertEqual(tb.get_status(tid3), states.FAILURE)
+        self.assertIsInstance(tb.get_result(tid3), KeyError)
+
+    def test_process_cleanup(self):
+        tb = CacheBackend(backend="memory://")
+        tb.process_cleanup()
+
+    def test_expires_as_int(self):
+        tb = CacheBackend(backend="memory://", expires=10)
+        self.assertEqual(tb.expires, 10)
+
+    def test_unknown_backend_raises_ImproperlyConfigured(self):
+        self.assertRaises(ImproperlyConfigured,
+                          CacheBackend, backend="unknown://")
+
+
+
+class test_get_best_memcache(unittest.TestCase):
+
+    def mock_memcache(self):
+        memcache = types.ModuleType("memcache")
+        memcache.Client = DummyClient
+        memcache.Client.__module__ = memcache.__name__
+        prev, sys.modules["memcache"] = sys.modules.get("memcache"), memcache
+        yield
+        if prev is not None:
+            sys.modules["memcache"] = prev
+        yield
+
+    def mock_pylibmc(self):
+        pylibmc = types.ModuleType("pylibmc")
+        pylibmc.Client = DummyClient
+        pylibmc.Client.__module__ = pylibmc.__name__
+        prev = sys.modules.get("pylibmc")
+        sys.modules["pylibmc"] = pylibmc
+        yield
+        if prev is not None:
+            sys.modules["pylibmc"] = prev
+        yield
+
+    def test_pylibmc(self):
+        pylibmc = self.mock_pylibmc()
+        pylibmc.next()
+        import __builtin__
+        sys.modules.pop("celery.backends.cache", None)
+        from celery.backends import cache
+        self.assertEqual(cache.get_best_memcache().__module__, "pylibmc")
+        pylibmc.next()
+
+    def test_memcache(self):
+
+        def with_no_pylibmc():
+            sys.modules.pop("celery.backends.cache", None)
+            from celery.backends import cache
+            self.assertEqual(cache.get_best_memcache().__module__, "memcache")
+
+        context = mask_modules("pylibmc")
+        context.__enter__()
+        try:
+            memcache = self.mock_memcache()
+            memcache.next()
+            with_no_pylibmc()
+            memcache.next()
+        finally:
+            context.__exit__(None, None, None)
+
+    def test_no_implementations(self):
+
+        def with_no_memcache_libs():
+            sys.modules.pop("celery.backends.cache", None)
+            from celery.backends import cache
+            self.assertRaises(ImproperlyConfigured, cache.get_best_memcache)
+
+        context = mask_modules("pylibmc", "memcache")
+        context.__enter__()
+        try:
+            with_no_memcache_libs()
+        finally:
+            context.__exit__(None, None, None)

celery/tests/test_backends/test_database.py

@@ -2,9 +2,13 @@
 import socket
 import unittest2 as unittest
 
+from datetime import datetime
+
 from celery.exceptions import ImproperlyConfigured
 
+from celery import conf
 from celery import states
+from celery.db.models import Task, TaskSet
 from celery.utils import gen_unique_id
 from celery.backends.database import DatabaseBackend
 
@@ -19,6 +23,17 @@ def __init__(self, data):
 
 class test_DatabaseBackend(unittest.TestCase):
 
+    def test_missing_dburi_raises_ImproperlyConfigured(self):
+        prev, conf.RESULT_DBURI = conf.RESULT_DBURI, None
+        try:
+            self.assertRaises(ImproperlyConfigured, DatabaseBackend)
+        finally:
+            conf.RESULT_DBURI = prev
+
+    def test_missing_task_id_is_PENDING(self):
+        tb = DatabaseBackend()
+        self.assertEqual(tb.get_status("xxx-does-not-exist"), states.PENDING)
+
     def test_mark_as_done(self):
         tb = DatabaseBackend()
 
@@ -84,3 +99,40 @@ def test_mark_as_failure(self):
     def test_process_cleanup(self):
         tb = DatabaseBackend()
         tb.process_cleanup()
+
+    def test_save___restore_taskset(self):
+        tb = DatabaseBackend()
+
+        tid = gen_unique_id()
+        res = {u"something": "special"}
+        self.assertEqual(tb.save_taskset(tid, res), res)
+
+        res2 = tb.restore_taskset(tid)
+        self.assertEqual(res2, res)
+
+        self.assertIsNone(tb.restore_taskset("xxx-nonexisting-id"))
+
+    def test_cleanup(self):
+        tb = DatabaseBackend()
+        for i in range(10):
+            tb.mark_as_done(gen_unique_id(), 42)
+            tb.save_taskset(gen_unique_id(), {"foo": "bar"})
+        s = tb.ResultSession()
+        for t in s.query(Task).all():
+            t.date_done = datetime.now() - tb.result_expires * 2
+        for t in s.query(TaskSet).all():
+            t.date_done = datetime.now() - tb.result_expires * 2
+        s.commit()
+        s.close()
+
+        tb.cleanup()
+        s2 = tb.ResultSession()
+        self.assertEqual(s2.query(Task).count(), 0)
+        self.assertEqual(s2.query(TaskSet).count(), 0)
+
+    def test_Task__repr__(self):
+        self.assertIn("foo", repr(Task("foo")))
+
+    def test_TaskSet__repr__(self):
+        self.assertIn("foo", repr(TaskSet("foo", None)))
+