Merge pull request #51 from bplotnick/be-more-defensive

Be more defensive about logging handlers being configured

py_zipkin/zipkin.py

@@ -174,6 +174,7 @@ def __init__(
         self.report_root_timestamp_override = report_root_timestamp
         self.use_128bit_trace_id = use_128bit_trace_id
         self.host = host
+        self.logging_configured = False
 
         # Spans that log a 'cs' timestamp can additionally record
         # 'sa' binary annotations that show where the request is going.
@@ -295,6 +296,7 @@ def start(self):
                 client_context=client_context
             )
             self.logging_context.start()
+            self.logging_configured = True
             return self
         else:
             # In the sampled case, patch the ZipkinLoggerHandler.
@@ -303,14 +305,22 @@ def start(self):
                 # the thread, multithreaded frameworks can get in strange states.
                 # The logging is not going to be correct in these cases, so we set
                 # a flag that turns off logging on __exit__.
-                if len(zipkin_logger.handlers) > 0:
-                    # Put span ID on logging handler. Assume there's only a single
-                    # handler, since all logging should be set up in this package.
-                    self.log_handler = zipkin_logger.handlers[0]
-                    # Store the old parent_span_id, probably None, in case we have
-                    # nested zipkin_spans
-                    self.old_parent_span_id = self.log_handler.parent_span_id
-                    self.log_handler.parent_span_id = self.zipkin_attrs.span_id
+                try:
+                    # Assume there's only a single handler, since all logging
+                    # should be set up in this package.
+                    log_handler = zipkin_logger.handlers[0]
+                except IndexError:
+                    return self
+                # Make sure it's not a NullHandler or something
+                if not isinstance(log_handler, ZipkinLoggerHandler):
+                    return self
+                # Put span ID on logging handler.
+                self.log_handler = zipkin_logger.handlers[0]
+                # Store the old parent_span_id, probably None, in case we have
+                # nested zipkin_spans
+                self.old_parent_span_id = self.log_handler.parent_span_id
+                self.log_handler.parent_span_id = self.zipkin_attrs.span_id
+                self.logging_configured = True
 
             return self
 
@@ -326,7 +336,7 @@ def stop(self, _exc_type=None, _exc_value=None, _exc_traceback=None):
         if self.do_pop_attrs:
             pop_zipkin_attrs()
 
-        if not self.zipkin_attrs or not self.zipkin_attrs.is_sampled:
+        if not self.logging_configured:
             return
 
         # Add the error annotation if an exception occurred

tests/zipkin_test.py

@@ -3,6 +3,7 @@
 
 import py_zipkin.zipkin as zipkin
 from py_zipkin.exception import ZipkinError
+from py_zipkin.logging_helper import null_handler
 from py_zipkin.logging_helper import ZipkinLoggerHandler
 from py_zipkin.thread_local import get_zipkin_attrs
 from py_zipkin.thrift import SERVER_ADDR_VAL
@@ -330,13 +331,15 @@ def test_span_context_no_zipkin_attrs(
     assert not push_zipkin_attrs_mock.called
 
 
+@pytest.mark.parametrize('handlers', [[], [null_handler]])
 @mock.patch('py_zipkin.thread_local._thread_local', autospec=True)
 @mock.patch('py_zipkin.zipkin.generate_random_64bit_string', autospec=True)
 @mock.patch('py_zipkin.zipkin.zipkin_logger', autospec=True)
 def test_span_context_sampled_no_handlers(
     zipkin_logger_mock,
     generate_string_mock,
     thread_local_mock,
+    handlers,
 ):
     zipkin_attrs = ZipkinAttrs(
         trace_id='1111111111111111',
@@ -347,14 +350,14 @@ def test_span_context_sampled_no_handlers(
     )
     thread_local_mock.zipkin_attrs = [zipkin_attrs]
 
-    zipkin_logger_mock.handlers = []
+    zipkin_logger_mock.handlers = handlers
     generate_string_mock.return_value = '1'
 
     context = zipkin.zipkin_span(
         service_name='my_service',
         port=5,
         transport_handler=mock.Mock(),
-        sample_rate=0.0,
+        sample_rate=None,
     )
     with context:
         # Assert that the new ZipkinAttrs were saved