def test_add_to_context(self):
mock_factory = mock.Mock(spec=ContextFactory)
mock_context = mock.Mock()
mock_span = mock.Mock(spec=Span)
observer = ContextObserver("some_attribute", mock_factory)
observer.on_root_span_created(mock_context, mock_span)
self.assertEqual(mock_context.some_attribute,
mock_factory.make_object_for_context.return_value)
def test_add_to_context(self):
mock_factory = mock.Mock(spec=ContextFactory)
mock_context = mock.Mock()
mock_span = mock.Mock(spec=Span)
observer = ContextObserver("some_attribute", mock_factory)
observer.on_server_span_created(mock_context, mock_span)
self.assertEqual(mock_context.some_attribute,
mock_factory.make_object_for_context.return_value)
def test_add_to_context_local(self):
mock_factory = mock.Mock(spec=ContextFactory)
mock_context = mock.Mock()
mock_local_span = mock.Mock(spec=LocalSpan)
mock_local_span.component_name = 'test_component'
mock_local_span.context = mock_context
observer = ContextObserver("some_attribute", mock_factory)
observer.on_child_span_created(mock_local_span)
self.assertEqual(mock_context.some_attribute,
mock_factory.make_object_for_context.return_value)
def test_add_to_context_local(self):
mock_factory = mock.Mock(spec=ContextFactory)
mock_context = mock.Mock()
mock_context = WrappedRequestContext(mock_context)
mock_local_span = mock.Mock(spec=LocalSpan)
mock_local_span.component_name = 'test_component'
mock_local_span.context = mock_context
observer = ContextObserver("some_attribute", mock_factory)
observer.on_child_span_created(mock_local_span)
self.assertEqual(mock_context.some_attribute,
mock_factory.make_object_for_context.return_value)
def add_to_context(self, name, context_factory):
"""Add an attribute to each request's context object.
On each request, the factory will be asked to create an appropriate
object to attach to the :term:`context object`.
:param str name: The attribute on the context object to attach the
created object to. This may also be used for metric/tracing
purposes so it should be descriptive.
:param baseplate.context.ContextFactory context_factory: A factory.
"""
# pylint: disable=cyclic-import
from baseplate.context import ContextObserver
self.register(ContextObserver(name, context_factory))