mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-08-02 16:50:24 +00:00
109 lines
2.9 KiB
Python
109 lines
2.9 KiB
Python
import pytest
|
|
|
|
from dbgpt.component import SystemApp
|
|
from dbgpt.util.tracer import (
|
|
DefaultTracer,
|
|
MemorySpanStorage,
|
|
Span,
|
|
SpanStorage,
|
|
SpanStorageType,
|
|
Tracer,
|
|
TracerManager,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def system_app():
|
|
return SystemApp()
|
|
|
|
|
|
@pytest.fixture
|
|
def storage(system_app: SystemApp):
|
|
ms = MemorySpanStorage(system_app)
|
|
system_app.register_instance(ms)
|
|
return ms
|
|
|
|
|
|
@pytest.fixture
|
|
def tracer(request, system_app: SystemApp):
|
|
if not request or not hasattr(request, "param"):
|
|
return DefaultTracer(system_app)
|
|
else:
|
|
span_storage_type = request.param.get(
|
|
"span_storage_type", SpanStorageType.ON_CREATE_END
|
|
)
|
|
return DefaultTracer(system_app, span_storage_type=span_storage_type)
|
|
|
|
|
|
@pytest.fixture
|
|
def tracer_manager(system_app: SystemApp, tracer: Tracer):
|
|
system_app.register_instance(tracer)
|
|
manager = TracerManager()
|
|
manager.initialize(system_app)
|
|
return manager
|
|
|
|
|
|
def test_start_and_end_span(tracer: Tracer):
|
|
span = tracer.start_span("operation")
|
|
assert isinstance(span, Span)
|
|
assert span.operation_name == "operation"
|
|
|
|
tracer.end_span(span)
|
|
assert span.end_time is not None
|
|
|
|
stored_span = tracer._get_current_storage().spans[0]
|
|
assert stored_span.span_id == span.span_id
|
|
|
|
|
|
def test_start_and_end_span_with_tracer_manager(tracer_manager: TracerManager):
|
|
span = tracer_manager.start_span("operation")
|
|
assert isinstance(span, Span)
|
|
assert span.operation_name == "operation"
|
|
|
|
tracer_manager.end_span(span)
|
|
assert span.end_time is not None
|
|
|
|
|
|
def test_parent_child_span_relation(tracer: Tracer):
|
|
parent_span = tracer.start_span("parent_operation")
|
|
child_span = tracer.start_span(
|
|
"child_operation", parent_span_id=parent_span.span_id
|
|
)
|
|
|
|
assert child_span.parent_span_id == parent_span.span_id
|
|
assert child_span.trace_id == parent_span.trace_id
|
|
|
|
tracer.end_span(child_span)
|
|
tracer.end_span(parent_span)
|
|
|
|
assert parent_span.operation_name in [
|
|
s.operation_name for s in tracer._get_current_storage().spans
|
|
]
|
|
assert child_span.operation_name in [
|
|
s.operation_name for s in tracer._get_current_storage().spans
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"tracer, expected_count, after_create_inc_count",
|
|
[
|
|
({"span_storage_type": SpanStorageType.ON_CREATE}, 1, 1),
|
|
({"span_storage_type": SpanStorageType.ON_END}, 1, 0),
|
|
({"span_storage_type": SpanStorageType.ON_CREATE_END}, 2, 1),
|
|
],
|
|
indirect=["tracer"],
|
|
)
|
|
def test_tracer_span_storage_type_and_with(
|
|
tracer: Tracer,
|
|
expected_count: int,
|
|
after_create_inc_count: int,
|
|
storage: SpanStorage,
|
|
):
|
|
span = tracer.start_span("new_span")
|
|
span.end()
|
|
assert len(storage.spans) == expected_count
|
|
|
|
with tracer.start_span("with_span") as ws:
|
|
assert len(storage.spans) == expected_count + after_create_inc_count
|
|
assert len(storage.spans) == expected_count + expected_count
|