Add optional error handler arg to register and add/remove functions (#34

)

* Add optional error handler arg to register and add/remove functions

* Update django dependency to use 0 instead of *

* Fix mocks

* Fix tests

* Fix call_count assertion

* Add tests for manual add/remove and unregister error handlers

* bump version

autocompleter/__init__.py

@@ -1,4 +1,4 @@
-VERSION = (1, 0, 2)
+VERSION = (1, 0, 3)
 
 from autocompleter.registry import registry, signal_registry
 from autocompleter.base import (

autocompleter/registry.py

@@ -1,6 +1,7 @@
 from collections import OrderedDict
 
 from django.db.models.signals import post_save, post_delete
+from functools import partial
 
 from autocompleter import settings
 
@@ -189,54 +190,96 @@ def del_ac_provider_setting(self, ac_name, provider, setting_name):
 registry = AutocompleterRegistry()
 
 
-def add_obj_to_autocompleter(sender, instance, created, **kwargs):
+def add_obj_to_autocompleter(sender, instance, created, add_error_handler = None, remove_error_handler = None, **kwargs):
     if instance is None:
         return
 
     provider_classes = registry.get_all_by_model(sender)
     for provider_class in provider_classes:
         provider = provider_class(instance)
         if provider.include_item():
-            provider.store()
+            try:
+                provider.store()
+            except Exception as e:
+                if add_error_handler:
+                    add_error_handler(instance, e)
+                else:
+                    raise e
         else:
             # If the item no longer passes the .include_item()
             # check then we need to remove it.
-            provider.remove()
+            try:
+                provider.remove()
+            except Exception as e:
+                if remove_error_handler:
+                    remove_error_handler(instance, e)
+                else:
+                    raise e
 
 
-def remove_obj_from_autocompleter(sender, instance, **kwargs):
+def remove_obj_from_autocompleter(sender, instance, remove_error_handler = None, **kwargs):
     if instance is None:
         return
 
     provider_classes = registry.get_all_by_model(sender)
     for provider_class in provider_classes:
-        provider_class(instance).remove()
+        try:
+            provider_class(instance).remove()
+        except Exception as e:
+            if remove_error_handler:
+                remove_error_handler(instance, e)
+            else:
+                raise e
+
+def remove_obj_from_autocompleter_with_error_handler(error_handler, sender, instance, **kwargs):
+    return remove_obj_from_autocompleter(sender, instance, remove_error_handler=error_handler, **kwargs)
 
+def add_obj_to_autocompleter_with_error_handler(add_error_handler, remove_error_handler, sender, instance, **kwargs):
+    return add_obj_to_autocompleter(sender, instance, add_error_handler=add_error_handler, remove_error_handler=remove_error_handler, **kwargs)
 
 class AutocompleterSignalRegistry(object):
-    def register(self, model):
+    DISPATCH_ID_FUNCTION_MAPPING = {}
+
+    def register(self, model, add_error_handler = None, remove_error_handler = None):
+
+        add_uid = "autocompleter.%s.add" % (model)
+        remove_uid = "autocompleter.%s.remove" % (model)
+
+        if add_error_handler or remove_error_handler:
+            remove_function = partial(remove_obj_from_autocompleter_with_error_handler, remove_error_handler)
+            add_function = partial(add_obj_to_autocompleter_with_error_handler, add_error_handler, remove_error_handler)
+
+            self.DISPATCH_ID_FUNCTION_MAPPING[add_uid] = add_function
+            self.DISPATCH_ID_FUNCTION_MAPPING[remove_uid] = remove_function
+
         post_save.connect(
-            add_obj_to_autocompleter,
+            self.DISPATCH_ID_FUNCTION_MAPPING.get(add_uid, add_obj_to_autocompleter),
             sender=model,
-            dispatch_uid="autocompleter.%s.add" % (model),
+            dispatch_uid=add_uid,
+
         )
         post_delete.connect(
-            remove_obj_from_autocompleter,
+            self.DISPATCH_ID_FUNCTION_MAPPING.get(remove_uid, remove_obj_from_autocompleter),
             sender=model,
-            dispatch_uid="autocompleter.%s.remove" % (model),
+            dispatch_uid=remove_uid,
         )
 
     def unregister(self, model):
+        add_uid = "autocompleter.%s.add" % (model)
+        remove_uid = "autocompleter.%s.remove" % (model)
         post_save.disconnect(
-            add_obj_to_autocompleter,
+            self.DISPATCH_ID_FUNCTION_MAPPING.get(add_uid, add_obj_to_autocompleter),
             sender=model,
-            dispatch_uid="autocompleter.%s.add" % (model),
+            dispatch_uid=add_uid,
         )
         post_delete.disconnect(
-            remove_obj_from_autocompleter,
+            self.DISPATCH_ID_FUNCTION_MAPPING.get(remove_uid, remove_obj_from_autocompleter),
             sender=model,
-            dispatch_uid="autocompleter.%s.remove" % (model),
+            dispatch_uid=remove_uid,
         )
 
+        self.DISPATCH_ID_FUNCTION_MAPPING.pop(add_uid, None)
+        self.DISPATCH_ID_FUNCTION_MAPPING.pop(remove_uid, None)
+
 
 signal_registry = AutocompleterSignalRegistry()

pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
     "Framework :: Django",
 ]
 dependencies = [
-    "Django >=3.2.*, <4.0",
+    "Django >=3.2.0, <4.0",
     "hiredis >= 1",
     "redis >= 3",
 ]

test_project/test_app/tests/test_storage.py

@@ -11,7 +11,9 @@
 )
 from test_app import calc_info
 from autocompleter import base, Autocompleter, registry, signal_registry
+from autocompleter.registry import add_obj_to_autocompleter, remove_obj_from_autocompleter
 from autocompleter import settings as auto_settings
+from unittest.mock import patch, MagicMock
 
 
 class StoringAndRemovingTestCase(AutocompleterTestCase):
@@ -482,6 +484,80 @@ def test_signal_based_add_and_remove(self):
 
         signal_registry.unregister(Stock)
 
+    @patch('autocompleter.base.AutocompleterProviderBase.store')
+    @patch('autocompleter.base.AutocompleterProviderBase.remove')
+    def test_signal_based_add_and_remove_error_handlers(self, mock_remove, mock_store):
+        """
+        Errors are properly handled when add or remove signal is sent.
+        """
+
+        remove_handler = MagicMock()
+        add_handler = MagicMock()
+
+        mock_remove.side_effect = Exception()
+        mock_store.side_effect = Exception()
+        aapl = Stock(symbol="AAPL", name="Apple", market_cap=50)
+
+        signal_registry.register(Stock, add_error_handler=add_handler, remove_error_handler=remove_handler)
+
+        aapl.save()
+        # There are 2 autocompleter providers for the Stock model. Therefore expect two error handler calls
+        self.assertEqual(add_handler.call_count, 2)
+
+        aapl.delete()
+        self.assertEqual(remove_handler.call_count, 2)
+
+        signal_registry.unregister(Stock)
+
+    @patch('autocompleter.base.AutocompleterProviderBase.store')
+    @patch('autocompleter.base.AutocompleterProviderBase.remove')
+    def test_add_and_remove_error_handlers(self, mock_remove, mock_store):
+        """
+        Errors are properly handled when adding or removing manually
+        """
+
+        remove_handler = MagicMock()
+        add_handler = MagicMock()
+
+        mock_remove.side_effect = Exception()
+        mock_store.side_effect = Exception()
+        aapl = Stock(symbol="AAPL", name="Apple", market_cap=50)
+
+        add_obj_to_autocompleter(Stock, aapl, False, add_error_handler=add_handler, remove_error_handler=remove_handler)
+        # There are 2 autocompleter providers for the Stock model. Therefore expect two error handler calls
+        self.assertEqual(add_handler.call_count, 2)
+
+        remove_obj_from_autocompleter(Stock, aapl, remove_error_handler=remove_handler)
+        self.assertEqual(remove_handler.call_count, 2)
+
+        signal_registry.unregister(Stock)
+
+    @patch('autocompleter.base.AutocompleterProviderBase.store')
+    @patch('autocompleter.base.AutocompleterProviderBase.remove')
+    def test_unregister_removes_error_handlers(self, mock_remove, mock_store):
+        """
+        Unregistering removes registered error handling
+        """
+        remove_handler = MagicMock()
+        add_handler = MagicMock()
+
+        mock_remove.side_effect = Exception()
+        mock_store.side_effect = Exception()
+        aapl = Stock(symbol="AAPL", name="Apple", market_cap=50)
+
+        signal_registry.register(Stock, add_error_handler=add_handler, remove_error_handler=remove_handler)
+
+        # re register Stock without error handlers
+        signal_registry.unregister(Stock)
+        signal_registry.register(Stock)
+
+        with self.assertRaises(Exception):
+            aapl.save()
+
+        with self.assertRaises(Exception):
+            aapl.delete()
+        signal_registry.unregister(Stock)
+
     def test_signal_based_update(self):
         """
         Turning on signals will automatically update objects in the autocompleter