mirror of
https://github.com/deepset-ai/haystack.git
synced 2025-07-06 16:41:15 +00:00
86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
|
import functools
|
||
|
|
import json
|
||
|
|
from typing import List, Dict
|
||
|
|
|
||
|
|
import ddtrace
|
||
|
|
import pytest
|
||
|
|
from _pytest.capture import CaptureFixture
|
||
|
|
from _pytest.monkeypatch import MonkeyPatch
|
||
|
|
|
||
|
|
from haystack.tracing.datadog import DatadogTracer
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture()
|
||
|
|
def datadog_tracer(monkeypatch: MonkeyPatch) -> ddtrace.Tracer:
|
||
|
|
# For the purpose of the tests we want to use the log writer
|
||
|
|
monkeypatch.setattr(ddtrace.Tracer, ddtrace.Tracer._use_log_writer.__name__, lambda *_: True)
|
||
|
|
|
||
|
|
tracer = ddtrace.Tracer()
|
||
|
|
|
||
|
|
return tracer
|
||
|
|
|
||
|
|
|
||
|
|
def get_traces_from_console(capfd: CaptureFixture) -> List[Dict]:
|
||
|
|
output = capfd.readouterr().out
|
||
|
|
parsed = json.loads(output)
|
||
|
|
nested_traces = parsed["traces"]
|
||
|
|
flattened = list(functools.reduce(lambda x, y: x + y, nested_traces, []))
|
||
|
|
|
||
|
|
return flattened
|
||
|
|
|
||
|
|
|
||
|
|
class TestDatadogTracer:
|
||
|
|
def test_opentelemetry_tracer(self, datadog_tracer: ddtrace.Tracer, capfd: CaptureFixture) -> None:
|
||
|
|
tracer = DatadogTracer(datadog_tracer)
|
||
|
|
|
||
|
|
with tracer.trace("test") as span:
|
||
|
|
span.set_tag("key", "value")
|
||
|
|
|
||
|
|
traces = get_traces_from_console(capfd)
|
||
|
|
assert len(traces) == 1
|
||
|
|
|
||
|
|
trace = traces[0]
|
||
|
|
|
||
|
|
assert trace["name"] == "test"
|
||
|
|
|
||
|
|
def test_tagging(self, datadog_tracer: ddtrace.Tracer, capfd: CaptureFixture) -> None:
|
||
|
|
tracer = DatadogTracer(datadog_tracer)
|
||
|
|
|
||
|
|
with tracer.trace("test", tags={"key1": "value1"}) as span:
|
||
|
|
span.set_tag("key2", "value2")
|
||
|
|
|
||
|
|
spans = get_traces_from_console(capfd)
|
||
|
|
assert len(spans) == 1
|
||
|
|
assert spans[0]["meta"]["key1"] == "value1"
|
||
|
|
assert spans[0]["meta"]["key2"] == "value2"
|
||
|
|
|
||
|
|
def test_current_span(self, datadog_tracer: ddtrace.Tracer, capfd: CaptureFixture) -> None:
|
||
|
|
tracer = DatadogTracer(datadog_tracer)
|
||
|
|
|
||
|
|
with tracer.trace("test"):
|
||
|
|
current_span = tracer.current_span()
|
||
|
|
assert tracer.current_span() is not None
|
||
|
|
|
||
|
|
current_span.set_tag("key1", "value1")
|
||
|
|
|
||
|
|
raw_span = current_span.raw_span()
|
||
|
|
assert raw_span is not None
|
||
|
|
assert isinstance(raw_span, ddtrace.Span)
|
||
|
|
|
||
|
|
raw_span.set_tag("key2", "value2")
|
||
|
|
|
||
|
|
spans = get_traces_from_console(capfd)
|
||
|
|
assert len(spans) == 1
|
||
|
|
assert spans[0]["meta"]["key1"] == "value1"
|
||
|
|
assert spans[0]["meta"]["key2"] == "value2"
|
||
|
|
|
||
|
|
def test_tracing_complex_values(self, datadog_tracer: ddtrace.Tracer, capfd: CaptureFixture) -> None:
|
||
|
|
tracer = DatadogTracer(datadog_tracer)
|
||
|
|
|
||
|
|
with tracer.trace("test") as span:
|
||
|
|
span.set_tag("key", {"a": 1, "b": [2, 3, 4]})
|
||
|
|
|
||
|
|
spans = get_traces_from_console(capfd)
|
||
|
|
assert len(spans) == 1
|
||
|
|
assert spans[0]["meta"]["key"] == '{"a": 1, "b": [2, 3, 4]}'
|