Improve test util to work as decorator or context processor

misago/acl/test.py

@@ -1,3 +1,4 @@
+from contextlib import ExitStack
 from functools import wraps
 from unittest.mock import patch
 
@@ -6,20 +7,31 @@ from .useracl import get_user_acl
 __all__ = ["patch_user_acl"]
 
 
-class patch_user_acl:
+class patch_user_acl(ExitStack):
     """Testing utility that patches get_user_acl results
 
-    Patch should be a dict or callable.
     Can be used as decorator or context manager.
+    
+    Accepts one or two arguments:
+    - patch_user_acl(acl_patch)
+    - patch_user_acl(user, acl_patch)
+
+    Patch should be a dict or callable.
     """
-    _global_patch = None
-    _user_patches = {}
 
-    def __init__(self, global_patch=None):
-        self._global_patch = global_patch
+    def __init__(self, *args):
+        super().__init__()
+
+        self._global_patch = None
+        self._user_patches = {}
 
-    def patch_user_acl(self, user, patch):
-        self._user_patches[user.id] = patch
+        if len(args) == 2:
+            user, patch = args
+            self._user_patches[user.id] = patch
+        elif len(args) == 1:
+            self._global_patch = args[0]
+        else:
+            raise ValueError("patch_user_acl takes one or two arguments.")
 
     def patched_get_user_acl(self, user, cache_versions):
         user_acl = get_user_acl(user, cache_versions)
@@ -40,10 +52,13 @@ class patch_user_acl:
             user_acl.update(acl_patch)
 
     def __enter__(self):
-        return self
-
-    def __exit__(self, *_):
-        self._user_patches = {}
+        super().__enter__()
+        self.enter_context(
+            patch(
+                "misago.acl.useracl.get_user_acl",
+                side_effect=self.patched_get_user_acl,
+            )
+        )
 
     def __call__(self, f):
         @wraps(f)
@@ -53,7 +68,6 @@ class patch_user_acl:
                     "misago.acl.useracl.get_user_acl",
                     side_effect=self.patched_get_user_acl,
                 ):
-                    new_args = args + (self.patch_user_acl,)
-                    return f(*new_args, **kwargs)
+                    return f(*args, **kwargs)
         
         return inner

misago/acl/tests/test_patching_user_acl.py

@@ -13,59 +13,61 @@ def callable_acl_patch(user, user_acl):
     user_acl["patched_for_user_id"] = user.id
 
 
-class PatchingUserACLInTestsTests(TestCase):
-    @patch_user_acl()
-    def test_decorator_adds_patching_function_to_test(self, patch_user_acl):
-        assert patch_user_acl
-
-    @patch_user_acl()
-    def test_patching_function_changes_user_permission_value(self, patch_user_acl):
-        user = User.objects.create_user("User", "user@example.com")
-        patch_user_acl(user, {"can_rename_users": 123})
-        user_acl = useracl.get_user_acl(user, cache_versions)
-        assert user_acl["can_rename_users"] == 123
-
-    @patch_user_acl()
-    def test_patching_function_adds_user_permission(self, patch_user_acl):
+class PatchingUserACLTests(TestCase):
+    @patch_user_acl({"is_patched": True})
+    def test_decorator_patches_all_users_acls_in_test(self):
         user = User.objects.create_user("User", "user@example.com")
-        patch_user_acl(user, {"new_user_permission": 123})
         user_acl = useracl.get_user_acl(user, cache_versions)
-        assert user_acl["new_user_permission"] == 123
+        assert user_acl["is_patched"]
 
-    def test_acl_patches_are_removed_after_test(self):
+    def test_decorator_removes_patches_after_test(self):
         user = User.objects.create_user("User", "user@example.com")
 
-        @patch_user_acl()
+        @patch_user_acl({"is_patched": True})
         def test_function(patch_user_acl):
-            patch_user_acl(user, {"is_patched": True})
             user_acl = useracl.get_user_acl(user, cache_versions)
             assert user_acl["is_patched"]
 
         user_acl = useracl.get_user_acl(user, cache_versions)
         assert "is_patched" not in user_acl
-    
-    @patch_user_acl({"is_patched": True})
-    def test_decorator_given_global_patch_patches_all_users_acls(self, _):
+
+    def test_context_manager_patches_all_users_acls_in_test(self):
         user = User.objects.create_user("User", "user@example.com")
-        user_acl = useracl.get_user_acl(user, cache_versions)
-        assert user_acl["is_patched"]
+        with patch_user_acl({"can_rename_users": "patched"}):
+            user_acl = useracl.get_user_acl(user, cache_versions)
+            assert user_acl["can_rename_users"] == "patched"
 
-    @patch_user_acl(callable_acl_patch)
-    def test_callable_global_patch_is_called_with_user_and_user_acl(self, _):
+    def test_context_manager_removes_patches_after_exit(self):
         user = User.objects.create_user("User", "user@example.com")
+
+        with patch_user_acl({"is_patched": True}):
+            user_acl = useracl.get_user_acl(user, cache_versions)
+            assert user_acl["is_patched"]
+
         user_acl = useracl.get_user_acl(user, cache_versions)
-        assert user_acl["patched_for_user_id"] == user.id
+        assert "is_patched" not in user_acl
 
-    @patch_user_acl({"is_patched": True})
-    def test_patching_function_overrides_global_patch(self, patch_user_acl):
+    def test_context_manager_patches_specified_user_acl(self):
         user = User.objects.create_user("User", "user@example.com")
-        patch_user_acl(user, {"is_patched": 123})
-        user_acl = useracl.get_user_acl(user, cache_versions)
-        assert user_acl["is_patched"] == 123
+        with patch_user_acl(user, {"can_rename_users": "patched"}):
+            user_acl = useracl.get_user_acl(user, cache_versions)
+            assert user_acl["can_rename_users"] == "patched"
+
+    def test_other_user_acl_is_not_changed_by_user_specific_context_manager(self):
+        patched_user = User.objects.create_user("User", "user@example.com")
+        other_user = User.objects.create_user("User2", "user2@example.com")
+        with patch_user_acl(patched_user, {"can_rename_users": "patched"}):
+            other_user_acl = useracl.get_user_acl(other_user, cache_versions)
+            assert other_user_acl["can_rename_users"] != "patched"
 
-    @patch_user_acl()
-    def test_callable_user_patch_is_called_with_user_and_user_acl(self, patch_user_acl):
+    @patch_user_acl(callable_acl_patch)
+    def test_callable_patch_is_called_with_user_and_acl_by_decorator(self):
         user = User.objects.create_user("User", "user@example.com")
-        patch_user_acl(user, callable_acl_patch)
         user_acl = useracl.get_user_acl(user, cache_versions)
         assert user_acl["patched_for_user_id"] == user.id
+
+    def test_callable_patch_is_called_with_user_and_acl_by_context_manager(self):
+        user = User.objects.create_user("User", "user@example.com")
+        with patch_user_acl(callable_acl_patch):
+            user_acl = useracl.get_user_acl(user, cache_versions)
+            assert user_acl["patched_for_user_id"] == user.id