refactor: The first refactored version for sdk release (#907)

Co-authored-by: chengfangyin2 <chengfangyin3@jd.com>

dbgpt/util/tracer/tests/test_base.py

@@ -0,0 +1,131 @@
+from typing import Dict
+from dbgpt.component import SystemApp
+
+from dbgpt.util.tracer import Span, SpanType, SpanStorage, Tracer
+
+
+# Mock implementations
+
+
+class MockSpanStorage(SpanStorage):
+    def __init__(self):
+        self.spans = []
+
+    def append_span(self, span: Span):
+        self.spans.append(span)
+
+
+class MockTracer(Tracer):
+    def __init__(self, system_app: SystemApp | None = None):
+        super().__init__(system_app)
+        self.current_span = None
+        self.storage = MockSpanStorage()
+
+    def append_span(self, span: Span):
+        self.storage.append_span(span)
+
+    def start_span(
+        self, operation_name: str, parent_span_id: str = None, metadata: Dict = None
+    ) -> Span:
+        trace_id = (
+            self._new_uuid() if parent_span_id is None else parent_span_id.split(":")[0]
+        )
+        span_id = f"{trace_id}:{self._new_uuid()}"
+        span = Span(
+            trace_id, span_id, SpanType.BASE, parent_span_id, operation_name, metadata
+        )
+        self.current_span = span
+        return span
+
+    def end_span(self, span: Span):
+        span.end()
+        self.append_span(span)
+
+    def get_current_span(self) -> Span:
+        return self.current_span
+
+    def _get_current_storage(self) -> SpanStorage:
+        return self.storage
+
+
+# Tests
+
+
+def test_span_creation():
+    span = Span(
+        "trace_id",
+        "span_id",
+        SpanType.BASE,
+        "parent_span_id",
+        "operation",
+        {"key": "value"},
+    )
+    assert span.trace_id == "trace_id"
+    assert span.span_id == "span_id"
+    assert span.parent_span_id == "parent_span_id"
+    assert span.operation_name == "operation"
+    assert span.metadata == {"key": "value"}
+
+
+def test_span_end():
+    span = Span("trace_id", "span_id")
+    assert span.end_time is None
+    span.end()
+    assert span.end_time is not None
+
+
+def test_mock_tracer_start_span():
+    tracer = MockTracer()
+    span = tracer.start_span("operation")
+    assert span.operation_name == "operation"
+    assert tracer.get_current_span() == span
+
+
+def test_mock_tracer_end_span():
+    tracer = MockTracer()
+    span = tracer.start_span("operation")
+    tracer.end_span(span)
+    assert span in tracer._get_current_storage().spans
+
+
+def test_mock_tracer_append_span():
+    tracer = MockTracer()
+    span = Span("trace_id", "span_id")
+    tracer.append_span(span)
+    assert span in tracer._get_current_storage().spans
+
+
+def test_parent_child_span_relation():
+    tracer = MockTracer()
+
+    # Start a parent span
+    parent_span = tracer.start_span("parent_operation")
+
+    # Start a child span with parent span's ID
+    child_span = tracer.start_span(
+        "child_operation", parent_span_id=parent_span.span_id
+    )
+
+    # Assert the relationships
+    assert child_span.parent_span_id == parent_span.span_id
+    assert (
+        child_span.trace_id == parent_span.trace_id
+    )  # Assuming children share the same trace ID
+
+    # End spans
+    tracer.end_span(child_span)
+    tracer.end_span(parent_span)
+
+    # Assert they are in the storage
+    assert child_span in tracer._get_current_storage().spans
+    assert parent_span in tracer._get_current_storage().spans
+
+
+# This test checks if unique UUIDs are being generated.
+# Note: This is a simple test and doesn't guarantee uniqueness for large numbers of UUIDs.
+
+
+def test_new_uuid_unique():
+    tracer = MockTracer()
+    uuid_set = {tracer._new_uuid() for _ in range(1000)}
+    assert len(uuid_set) == 1000