diff --git a/tests/unit/sensors/test_distributed_tracing.py b/tests/unit/sensors/test_distributed_tracing.py new file mode 100644 index 000000000..2ae828a65 --- /dev/null +++ b/tests/unit/sensors/test_distributed_tracing.py @@ -0,0 +1,273 @@ +from unittest.mock import MagicMock, Mock + +import opentracing +import pytest + +from faust.sensors.distributed_tracing import TracingSensor +from faust.types import TP + +TP1 = TP("topic", 3) + + +@pytest.fixture() +def tracer(*, monkeypatch): + # ``app_tracer``/``kafka_tracer`` both return the module-global + # ``opentracing.tracer``; swap it for a mock so we can assert on the + # spans the sensor creates. + tracer = MagicMock(name="tracer") + monkeypatch.setattr("opentracing.tracer", tracer) + return tracer + + +@pytest.fixture() +def sensor(): + return TracingSensor() + + +def _message(*, headers=None, topic="topic", key=b"key"): + message = Mock(name="message") + message.headers = headers if headers is not None else [] + message.topic = topic + message.key = key + message.span = None + # ``stream_meta`` is only set on the message by the sensor; start absent + # so the getattr(..., None) branch is exercised. + del message.stream_meta + return message + + +class Test_TracingSensor: + def test_app_and_kafka_tracer(self, *, sensor, tracer): + assert sensor.app_tracer is tracer + assert sensor.kafka_tracer is tracer + + @pytest.mark.asyncio + async def test_stop__no_sessions(self, *, sensor): + # aiohttp_sessions defaults to None -> just delegates to super().stop(). + assert sensor.aiohttp_sessions is None + await sensor.stop() + + @pytest.mark.asyncio + async def test_stop__closes_sessions(self, *, sensor): + session = Mock(name="session") + session.close = Mock(return_value=_awaitable()) + sensor.aiohttp_sessions = {"a": session} + await sensor.stop() + session.close.assert_called_once_with() + + def test_on_message_in__no_headers(self, *, sensor, tracer, monkeypatch): + set_current_span = Mock(name="set_current_span") + monkeypatch.setattr( + "faust.sensors.distributed_tracing.set_current_span", set_current_span + ) + span = tracer.start_span.return_value + message = _message(headers=[]) + + sensor.on_message_in(TP1, 0, message) + + tracer.extract.assert_not_called() + tracer.start_span.assert_called_once_with(operation_name="consume-from-topic") + set_current_span.assert_called_once_with(span) + span.set_tag.assert_any_call("kafka-topic", TP1.topic) + span.set_tag.assert_any_call("kafka-partition", TP1.partition) + span.set_tag.assert_any_call("kafka-key", message.key) + span.__enter__.assert_called_once_with() + assert message.span is span + + def test_on_message_in__with_headers(self, *, sensor, tracer, monkeypatch): + monkeypatch.setattr( + "faust.sensors.distributed_tracing.set_current_span", Mock() + ) + follows_from = Mock(name="follows_from") + monkeypatch.setattr("opentracing.follows_from", follows_from) + message = _message(headers=[(b"uber-trace-id", b"123")]) + + sensor.on_message_in(TP1, 0, message) + + tracer.extract.assert_called_once() + follows_from.assert_called_once_with(tracer.extract.return_value) + tracer.start_span.assert_called_once_with( + operation_name="consume-from-topic", + references=follows_from.return_value, + ) + + def test_on_stream_event_in__no_parent_span(self, *, sensor, monkeypatch): + start_child = Mock(name="start_child_span") + monkeypatch.setattr("opentracing.start_child_span", start_child) + message = _message() + message.span = None + event = Mock(name="event", message=message) + + sensor.on_stream_event_in(TP1, 0, Mock(name="stream"), event) + + start_child.assert_not_called() + # stream_meta is created even when there is no parent span. + assert message.stream_meta == {} + + def test_on_stream_event_in__with_parent_span(self, *, sensor, monkeypatch): + stream_span = MagicMock(name="stream_span") + start_child = Mock(name="start_child_span", return_value=stream_span) + monkeypatch.setattr("opentracing.start_child_span", start_child) + parent_span = Mock(name="parent_span") + message = _message() + message.span = parent_span + event = Mock(name="event", message=message) + stream = Mock(name="stream", concurrency_index=2, prefix="p") + + sensor.on_stream_event_in(TP1, 0, stream, event) + + start_child.assert_called_once_with(parent_span, "job-topic") + stream_span.set_tag.assert_any_call("stream-concurrency-index", 2) + stream_span.set_tag.assert_any_call("stream-prefix", "p") + stream_span.__enter__.assert_called_once_with() + assert message.stream_meta["stream_spans"][stream] is stream_span + + def test_on_stream_event_in__reuses_existing_meta(self, *, sensor, monkeypatch): + # stream_meta and its stream_spans dict already exist -> the sensor + # must reuse them rather than recreate (covers the not-None branches). + stream_span = MagicMock(name="stream_span") + monkeypatch.setattr( + "opentracing.start_child_span", + Mock(return_value=stream_span), + ) + other_stream, other_span = Mock(name="other"), Mock(name="other_span") + message = _message() + message.span = Mock(name="parent_span") + message.stream_meta = {"stream_spans": {other_stream: other_span}} + event = Mock(name="event", message=message) + stream = Mock(name="stream", concurrency_index=0, prefix="p") + + sensor.on_stream_event_in(TP1, 0, stream, event) + + spans = message.stream_meta["stream_spans"] + assert spans[other_stream] is other_span # untouched + assert spans[stream] is stream_span # added to the same dict + + def test_on_stream_event_out__finishes_span(self, *, sensor): + span = Mock(name="span") + stream = Mock(name="stream") + message = _message() + message.stream_meta = {"stream_spans": {stream: span}} + event = Mock(name="event", message=message) + + sensor.on_stream_event_out(TP1, 0, stream, event) + + span.finish.assert_called_once_with() + assert stream not in message.stream_meta["stream_spans"] + + def test_on_stream_event_out__no_spans(self, *, sensor): + message = _message() + event = Mock(name="event", message=message) + # No stream_meta -> created empty, nothing to finish (no error). + sensor.on_stream_event_out(TP1, 0, Mock(name="stream"), event) + assert message.stream_meta == {} + + def test_on_stream_event_out__stream_not_present(self, *, sensor): + message = _message() + message.stream_meta = {"stream_spans": {Mock(name="other"): Mock()}} + event = Mock(name="event", message=message) + # pop returns None -> no finish, no crash. + sensor.on_stream_event_out(TP1, 0, Mock(name="stream"), event) + + def test_on_message_out__with_span(self, *, sensor): + message = _message() + message.span = Mock(name="span") + sensor.on_message_out(TP1, 0, message) + message.span.finish.assert_called_once_with() + + def test_on_message_out__no_span(self, *, sensor): + message = _message() + message.span = None + sensor.on_message_out(TP1, 0, message) # no crash + + def test_on_send_initiated__with_parent(self, *, sensor, monkeypatch): + span = MagicMock(name="span") + start_child = Mock(name="start_child_span", return_value=span) + monkeypatch.setattr("opentracing.start_child_span", start_child) + parent = Mock(name="parent_span") + monkeypatch.setattr( + "faust.sensors.distributed_tracing.current_span", + Mock(return_value=parent), + ) + inject = Mock(name="trace_inject_headers") + sensor.trace_inject_headers = inject + message = Mock(name="pending", headers={"h": "1"}) + + ret = sensor.on_send_initiated(Mock(), "topic", message, 1, 2) + + start_child.assert_called_once_with(parent, "produce-to-topic") + span.set_tag.assert_called_once_with("kafka-headers", {"h": "1"}) + inject.assert_called_once_with(span, message.headers) + span.__enter__.assert_called_once_with() + assert ret == {"span": span} + + def test_on_send_initiated__no_parent(self, *, sensor, monkeypatch): + monkeypatch.setattr( + "faust.sensors.distributed_tracing.current_span", + Mock(return_value=None), + ) + message = Mock(name="pending", headers=None) + ret = sensor.on_send_initiated(Mock(), "topic", message, 1, 2) + assert ret == {"span": None} + + def test_on_send_completed__with_span(self, *, sensor): + span = Mock(name="span") + metadata = Mock(topic="t", partition=1, offset=42) + sensor.on_send_completed(Mock(), {"span": span}, metadata) + span.set_tag.assert_any_call("kafka-topic", "t") + span.set_tag.assert_any_call("kafka-partition", 1) + span.set_tag.assert_any_call("kafka-offset", 42) + span.finish.assert_called_once_with() + + def test_on_send_completed__no_span(self, *, sensor): + sensor.on_send_completed(Mock(), {"span": None}, Mock()) # no crash + + def test_on_send_error__with_span(self, *, sensor): + span = Mock(name="span") + exc = ValueError("boom") + sensor.on_send_error(Mock(), exc, {"span": span}) + span.set_tag.assert_called_once_with(opentracing.ext.tags.ERROR, "true") + span.log_kv.assert_called_once() + span.finish.assert_called_once_with(exception=exc) + + def test_on_send_error__no_span(self, *, sensor): + sensor.on_send_error(Mock(), ValueError(), {"span": None}) # no crash + + def test_trace_inject_headers__success(self, *, sensor, tracer, monkeypatch): + merge_headers = Mock(name="merge_headers") + monkeypatch.setattr( + "faust.sensors.distributed_tracing.merge_headers", merge_headers + ) + span = Mock(name="span") + headers = {"a": "1"} + + ret = sensor.trace_inject_headers(span, headers) + + tracer.inject.assert_called_once() + merge_headers.assert_called_once() + assert ret is headers + + def test_trace_inject_headers__no_tracer(self, *, sensor, monkeypatch): + # app_tracer is None -> inject skipped, headers returned unchanged. + sensor.__dict__["app_tracer"] = None + headers = {"a": "1"} + assert sensor.trace_inject_headers(Mock(), headers) is headers + + def test_trace_inject_headers__no_span(self, *, sensor, tracer): + headers = {"a": "1"} + assert sensor.trace_inject_headers(None, headers) is headers + tracer.inject.assert_not_called() + + def test_trace_inject_headers__exception_returns_none(self, *, sensor, tracer): + tracer.inject.side_effect = RuntimeError("nope") + assert sensor.trace_inject_headers(Mock(name="span"), {"a": "1"}) is None + + def test_on_threaded_producer_buffer_processed(self, *, sensor): + assert sensor.on_threaded_producer_buffer_processed(Mock(), 10) is None + + +def _awaitable(result=None): + async def _coro(): + return result + + return _coro()