feat(signals) Add a context decorator to disable the signals

Author: Clement Denoix

algoliasearch_django/decorators.py

@@ -1,3 +1,39 @@
+try:
+    # ContextDecorator was introduced in Python 3.2
+    from contextlib import ContextDecorator
+except ImportError:
+    ContextDecorator = None
+from functools import WRAPPER_ASSIGNMENTS, wraps
+
+from django.db.models.signals import post_save, pre_delete
+
+from . import algolia_engine
+
+
+def available_attrs(fn):
+    """
+    Return the list of functools-wrappable attributes on a callable.
+    This was required as a workaround for http://bugs.python.org/issue3445
+    under Python 2.
+    """
+    return WRAPPER_ASSIGNMENTS
+
+
+if ContextDecorator is None:
+    # ContextDecorator was introduced in Python 3.2
+    # See https://docs.python.org/3/library/contextlib.html#contextlib.ContextDecorator
+    class ContextDecorator:
+        """
+        A base class that enables a context manager to also be used as a decorator.
+        """
+        def __call__(self, func):
+            @wraps(func, assigned=available_attrs(func))
+            def inner(*args, **kwargs):
+                with self:
+                    return func(*args, **kwargs)
+            return inner
+
+
 def register(model):
     """
     Register the given model class and wrapped AlgoliaIndex class with the Algolia engine:
@@ -17,3 +53,44 @@ def _algolia_engine_wrapper(index_class):
 
         return index_class
     return _algolia_engine_wrapper
+
+
+class disable_auto_indexing(ContextDecorator):
+    """
+    A context decorator to disable the auto-indexing behaviour of the AlgoliaIndex
+
+    Can be used either as a context manager or a method decorator:
+    >>> with disable_auto_indexing():
+    >>>     my_object.save()
+
+    >>> @disable_auto_indexing()
+    >>> big_operation()
+    """
+
+    def __init__(self, model=None):
+        if model is not None:
+            self.models = [model]
+        else:
+            self.models = algolia_engine._AlgoliaEngine__registered_models
+
+    def __enter__(self):
+        for model in self.models:
+            post_save.disconnect(
+                algolia_engine._AlgoliaEngine__post_save_receiver,
+                sender=model
+            )
+            pre_delete.disconnect(
+                algolia_engine._AlgoliaEngine__pre_delete_receiver,
+                sender=model
+            )
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        for model in self.models:
+            post_save.connect(
+                algolia_engine._AlgoliaEngine__post_save_receiver,
+                sender=model
+            )
+            pre_delete.connect(
+                algolia_engine._AlgoliaEngine__pre_delete_receiver,
+                sender=model
+            )

tests/test_decorators.py

@@ -0,0 +1,108 @@
+from mock import (
+    ANY,
+    call,
+    patch
+)
+
+from django.test import TestCase
+
+from algoliasearch_django import algolia_engine
+from algoliasearch_django.decorators import disable_auto_indexing
+
+from .factories import UserFactory, WebsiteFactory
+from .models import User
+
+
+class DecoratorsTestCase(TestCase):
+    def test_disable_auto_indexing_as_decorator_for_all(self):
+        """Test that the `disable_auto_indexing` should work as a decorator for all the model"""
+
+        @disable_auto_indexing()
+        def decorated_operation():
+            WebsiteFactory()
+            UserFactory()
+
+        def non_decorated_operation():
+            WebsiteFactory()
+            UserFactory()
+
+        with patch.object(algolia_engine, 'save_record') as mocked_save_record:
+            decorated_operation()
+
+        # The decorated method should have prevented the indexing operations
+        mocked_save_record.assert_not_called()
+
+        with patch.object(algolia_engine, 'save_record') as mocked_save_record:
+            non_decorated_operation()
+
+        # The non-decorated method is not preventing the indexing operations
+        # (the signal was correctly re-connected for both of the models)
+        mocked_save_record.assert_has_calls([
+            call(
+                ANY,
+                created=True,
+                raw=False,
+                sender=ANY,
+                signal=ANY,
+                update_fields=None,
+                using=ANY
+            )
+        ] * 2)
+
+    def test_disable_auto_indexing_as_decorator_for_model(self):
+        """Test that the `disable_auto_indexing` should work as a decorator for a specific model"""
+
+        @disable_auto_indexing(model=User)
+        def decorated_operation():
+            WebsiteFactory()
+            UserFactory()
+
+        def non_decorated_operation():
+            WebsiteFactory()
+            UserFactory()
+
+        with patch.object(algolia_engine, 'save_record') as mocked_save_record:
+            decorated_operation()
+
+        # The decorated method should have prevented the indexing operation for the `User` model
+        # but not for the `Website` model (we get only one call)
+        mocked_save_record.assert_called_once_with(
+            ANY,
+            created=True,
+            raw=False,
+            sender=ANY,
+            signal=ANY,
+            update_fields=None,
+            using=ANY
+        )
+
+        with patch.object(algolia_engine, 'save_record') as mocked_save_record:
+            non_decorated_operation()
+
+        # The non-decorated method is not preventing the indexing operations
+        # (the signal was correctly re-connected for both of the models)
+        mocked_save_record.assert_has_calls([
+            call(
+                ANY,
+                created=True,
+                raw=False,
+                sender=ANY,
+                signal=ANY,
+                update_fields=None,
+                using=ANY
+            )
+        ] * 2)
+
+    def test_disable_auto_indexing_as_context_manager(self):
+        """Test that the `disable_auto_indexing` should work as a context manager"""
+
+        with patch.object(algolia_engine, 'save_record') as mocked_save_record:
+            with disable_auto_indexing():
+                WebsiteFactory()
+
+        mocked_save_record.assert_not_called()
+
+        with patch.object(algolia_engine, 'save_record') as mocked_save_record:
+            WebsiteFactory()
+
+        mocked_save_record.assert_called_once()