1 files changed, 24 insertions, 0 deletions
diff --git a/t/mocks.py b/t/mocks.py
new file mode 100644
index 0000000..f91ae9b
--- /dev/null
+++ b/t/mocks.py
@@ -0,0 +1,24 @@
+from unittest.mock import Mock
+
+class _ContextMock(Mock):
+    """Dummy class implementing __enter__ and __exit__
+    as the :keyword:`with` statement requires these to be implemented
+    in the class, not just the instance."""
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc_info):
+        pass
+
+
+def ContextMock(*args, **kwargs):
+    """Mock that mocks :keyword:`with` statement contexts."""
+    obj = _ContextMock(*args, **kwargs)
+    obj.attach_mock(_ContextMock(), '__enter__')
+    obj.attach_mock(_ContextMock(), '__exit__')
+    obj.__enter__.return_value = obj
+    # if __exit__ return a value the exception is ignored,
+    # so it must return None here.
+    obj.__exit__.return_value = None
+    return obj