浏览代码

celery.tests.utils.Mock: With attrs argument to set arbitrary attributes on the mock object

Ask Solem 14 年之前
父节点
当前提交
7d171f8ab8
共有 1 个文件被更改,包括 10 次插入0 次删除
  1. 10 0
      celery/tests/utils.py

+ 10 - 0
celery/tests/utils.py

@@ -23,6 +23,7 @@ except ImportError:
     from celery.tests.utils import fallback_contextmanager
     contextmanager = fallback_contextmanager  # noqa
 
+import mock
 
 from nose import SkipTest
 
@@ -31,6 +32,15 @@ from celery.utils import noop
 from celery.utils.functional import wraps
 
 
+class Mock(mock.Mock):
+
+    def __init__(self, *args, **kwargs):
+        attrs = kwargs.pop("attrs", None) or {}
+        super(Mock, self).__init__(*args, **kwargs)
+        for attr_name, attr_value in attrs.items():
+            setattr(self, attr_name, attr_value)
+
+
 class AppCase(unittest.TestCase):
 
     def setUp(self):