feat: add support to start/stop on lock test helper

Author: joaodaher

drf_kit/tests.py

@@ -129,6 +129,19 @@ def __init__(self):
                 self.call_args = []
                 self.call_kwargs = {}
 
+            def start(self):
+                setattr(ConnectionProxy, "lock", mocked_lock)
+
+            def stop(self):
+                delattr(ConnectionProxy, "lock")
+
+            def __enter__(self):
+                self.start()
+                return self
+
+            def __exit__(self, exc_type, exc_val, exc_tb):
+                self.stop()
+
             def assert_called(self):
                 msg = "Expected to be called, but it was not"
                 assert self.call_count > 0, msg
@@ -174,10 +187,4 @@ def mocked_lock(*lock_args, **lock_kwargs):
 
             yield
 
-        @contextmanager
-        def lock_assertion():
-            setattr(ConnectionProxy, "lock", mocked_lock)
-            yield assertion
-            delattr(ConnectionProxy, "lock")
-
-        return lock_assertion()
+        return assertion

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "drf-kit"
-version = "1.8.0"
+version = "1.8.1"
 description = "DRF Toolkit"
 authors = ["eduK <[email protected]>"]
 packages = [

test_app/tests/tests_test_helpers.py

@@ -95,3 +95,12 @@ def _effect():
                 tasks.LockableTask().run()
 
         lock.assert_called()
+
+    def test_use_with_start(self):
+        lock = self.patch_cache_lock()
+        lock.start()
+
+        tasks.LockableTask().run()
+        lock.assert_called()
+
+        lock.stop()