pyinfra/tests/unit_test/opentelemetry_test.py
2024-01-24 13:26:01 +01:00

58 lines
2.0 KiB
Python

import json
from time import sleep
from opentelemetry.instrumentation.pika import PikaInstrumentor
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter, SpanExportResult
from opentelemetry import trace
class MySpanExporter(SpanExporter):
def __init__(self):
self.traces = []
def export(self, spans):
for span in spans:
self.traces.append(span.to_json())
return SpanExportResult.SUCCESS
def shutdown(self):
pass
class TestOpenTelemetry:
def test_queue_messages_are_traced(self, queue_manager, input_message, stop_message):
service_name = "deine-mutter-serivce"
resource = Resource(attributes={"service.name": service_name})
provider = TracerProvider(resource=resource)
# processor = BatchSpanProcessor(ConsoleSpanExporter())
exporter = MySpanExporter()
processor = BatchSpanProcessor(exporter)
provider.add_span_processor(processor)
# otlp_exporter = OTLPSpanExporter(endpoint=self.tracing_endpoint)
# span_processor = BatchSpanProcessor(otlp_exporter)
# trace.get_tracer_provider().add_span_processor(span_processor)
# Sets the global default tracer provider
trace.set_tracer_provider(provider)
# Creates a tracer from the global tracer provider
tracer = trace.get_tracer("my.tracer.name")
PikaInstrumentor().instrument()
queue_manager.purge_queues()
queue_manager.publish_message_to_input_queue(input_message)
queue_manager.publish_message_to_input_queue(stop_message)
def callback(_):
sleep(2)
return {"flat": "earth"}
queue_manager.start_consuming(callback)
for exported_trace in exporter.traces:
exported_trace = json.loads(exported_trace)
assert exported_trace["resource"]["attributes"]["service.name"] == service_name