allow subclasses of SentryHandler to be modified by the celery logging signal, resolves #961

2 files changed, 124 insertions, 2 deletions

diff --git a/raven/contrib/celery/__init__.py b/raven/contrib/celery/__init__.py
index d4fd759..02f3ab0 100644
--- a/raven/contrib/celery/__init__.py
+++ b/raven/contrib/celery/__init__.py
@@ -43,7 +43,7 @@ def register_logger_signal(client, logger=None, loglevel=logging.ERROR):
         # that the CeleryFilter is installed.
         # If one is found, we do not attempt to install another one.
         for h in logger.handlers:
-            if type(h) == SentryHandler:
+            if isinstance(h, SentryHandler):
                 h.addFilter(filter_)
                 return False
 
diff --git a/tests/contrib/test_celery.py b/tests/contrib/test_celery.py
index 95a3b74..cfa5b0f 100644
--- a/tests/contrib/test_celery.py
+++ b/tests/contrib/test_celery.py
@@ -1,8 +1,14 @@
 from __future__ import absolute_import
 
 import celery
+import logging
 
-from raven.contrib.celery import SentryCeleryHandler
+from raven.contrib.celery import (
+    SentryCeleryHandler,
+    register_logger_signal,
+    CeleryFilter
+)
+from raven.handlers.logging import SentryHandler
 from raven.utils.testutils import InMemoryClient, TestCase
 
 
@@ -38,3 +44,119 @@ class CeleryTestCase(TestCase):
         dummy_task.delay(1, 2)
         dummy_task.delay(1, 0)
         assert len(self.client.events) == 0
+
+
+class CeleryLoggingHandlerTestCase(TestCase):
+    def setUp(self):
+        super(CeleryLoggingHandlerTestCase, self).setUp()
+
+        self.client = InMemoryClient()
+
+        # register the logger signal
+        # and unregister the signal when the test is done
+        register_logger_signal(self.client)
+        receiver = celery.signals.after_setup_logger.receivers[0][1]
+        self.addCleanup(celery.signals.after_setup_logger.disconnect, receiver)
+
+        # remove any existing handlers and restore
+        # them when complete
+        self.root = logging.getLogger()
+        for handler in self.root.handlers:
+            self.root.removeHandler(handler)
+            self.addCleanup(self.root.addHandler, handler)
+
+    def test_handler_added(self):
+        # Given: there are no handlers configured
+        assert not self.root.handlers
+
+        # When: the after_setup_logger signal is sent
+        celery.signals.after_setup_logger.send(
+            sender=None, logger=self.root,
+            loglevel=logging.WARNING, logfile=None,
+            format=u'', colorize=False,
+        )
+
+        # Then: there is 1 new handler
+        assert len(self.root.handlers) == 1
+
+        # Then: the new handler is an instance of
+        # `raven.handlers.logging.SentryHandler`
+        handler = self.root.handlers[0]
+        assert isinstance(handler, SentryHandler)
+
+        # Then: the handler has 1 filter
+        assert len(handler.filters) == 1
+
+        # Then: the filter is a CeleryFilter
+        _filter = handler.filters[0]
+        assert isinstance(_filter, CeleryFilter)
+
+        # set up the handler to be removed once the test is done
+        self.addCleanup(self.root.removeHandler, handler)
+
+    def test_handler_updated(self):
+
+        # Given: there is 1 preconfigured SentryHandler
+        # with no filters
+        handler = SentryHandler(self.client)
+        assert not handler.filters
+        self.root.addHandler(handler)
+        # set up the handler to be removed once the test is done
+        self.addCleanup(self.root.removeHandler, handler)
+
+        # When: the after_setup_logger signal is sent
+        celery.signals.after_setup_logger.send(
+            sender=None, logger=self.root,
+            loglevel=logging.WARNING, logfile=None,
+            format=u'', colorize=False,
+        )
+
+        # Then: there is still just 1 handler
+        assert len(self.root.handlers) == 1
+
+        # Then: the existing handler is an instance of
+        # `raven.handlers.logging.SentryHandler`
+        handler = self.root.handlers[0]
+        assert isinstance(handler, SentryHandler)
+
+        # Then: the existing handler has 1 filter
+        assert len(handler.filters) == 1
+
+        # Then: the filter is a CeleryFilter
+        _filter = handler.filters[0]
+        assert isinstance(_filter, CeleryFilter)
+
+    def test_subclassed_handler_updated(self):
+
+        # Given: there is 1 preconfigured CustomHandler
+        # with no filters
+        class CustomHandler(SentryHandler):
+            pass
+
+        handler = CustomHandler(self.client)
+        assert not handler.filters
+        self.root.addHandler(handler)
+        # set up the handler to be removed once the test is done
+        self.addCleanup(self.root.removeHandler, handler)
+
+        # When: the after_setup_logger signal is sent
+        celery.signals.after_setup_logger.send(
+            sender=None, logger=self.root,
+            loglevel=logging.WARNING, logfile=None,
+            format=u'', colorize=False,
+        )
+
+        # Then: there is still just 1 handler
+        assert len(self.root.handlers) == 1
+
+        # Then: the existing handler is an instance of
+        # `CustomHandler`
+        handler = self.root.handlers[0]
+        assert isinstance(handler, CustomHandler)
+
+        # Then: the existing handler has 1 filter
+        assert len(handler.filters) == 1
+
+        # Then: the filter is a CeleryFilter
+        _filter = handler.filters[0]
+        assert isinstance(_filter, CeleryFilter)