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

dbgpt/util/tracer/tests/test_span_storage.py

@@ -0,0 +1,174 @@
+import os
+import pytest
+import asyncio
+import json
+import tempfile
+import time
+from unittest.mock import patch
+from datetime import datetime, timedelta
+
+from dbgpt.util.tracer import (
+    SpanStorage,
+    FileSpanStorage,
+    Span,
+    SpanType,
+    SpanStorageContainer,
+)
+
+
+@pytest.fixture
+def storage(request):
+    if not request or not hasattr(request, "param"):
+        file_does_not_exist = False
+    else:
+        file_does_not_exist = request.param.get("file_does_not_exist", False)
+
+    if file_does_not_exist:
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            filename = os.path.join(tmp_dir, "non_existent_file.jsonl")
+            storage_instance = FileSpanStorage(filename)
+            yield storage_instance
+    else:
+        with tempfile.NamedTemporaryFile(delete=True) as tmp_file:
+            filename = tmp_file.name
+            storage_instance = FileSpanStorage(filename)
+            yield storage_instance
+
+
+@pytest.fixture
+def storage_container(request):
+    if not request or not hasattr(request, "param"):
+        batch_size = 10
+        flush_interval = 10
+    else:
+        batch_size = request.param.get("batch_size", 10)
+        flush_interval = request.param.get("flush_interval", 10)
+    storage_container = SpanStorageContainer(
+        batch_size=batch_size, flush_interval=flush_interval
+    )
+    yield storage_container
+
+
+def read_spans_from_file(filename):
+    with open(filename, "r") as f:
+        return [json.loads(line) for line in f.readlines()]
+
+
+def test_write_span(storage: SpanStorage):
+    span = Span("1", "a", SpanType.BASE, "b", "op1")
+    storage.append_span(span)
+    time.sleep(0.1)
+
+    spans_in_file = read_spans_from_file(storage.filename)
+    assert len(spans_in_file) == 1
+    assert spans_in_file[0]["trace_id"] == "1"
+
+
+def test_incremental_write(storage: SpanStorage):
+    span1 = Span("1", "a", SpanType.BASE, "b", "op1")
+    span2 = Span("2", "c", SpanType.BASE, "d", "op2")
+
+    storage.append_span(span1)
+    storage.append_span(span2)
+    time.sleep(0.1)
+
+    spans_in_file = read_spans_from_file(storage.filename)
+    assert len(spans_in_file) == 2
+
+
+def test_sync_and_async_append(storage: SpanStorage):
+    span = Span("1", "a", SpanType.BASE, "b", "op1")
+
+    storage.append_span(span)
+
+    async def async_append():
+        storage.append_span(span)
+
+    asyncio.run(async_append())
+
+    time.sleep(0.1)
+    spans_in_file = read_spans_from_file(storage.filename)
+    assert len(spans_in_file) == 2
+
+
+@pytest.mark.parametrize("storage", [{"file_does_not_exist": True}], indirect=True)
+def test_non_existent_file(storage: SpanStorage):
+    span = Span("1", "a", SpanType.BASE, "b", "op1")
+    span2 = Span("2", "c", SpanType.BASE, "d", "op2")
+    storage.append_span(span)
+    time.sleep(0.1)
+
+    spans_in_file = read_spans_from_file(storage.filename)
+    assert len(spans_in_file) == 1
+
+    storage.append_span(span2)
+    time.sleep(0.1)
+    spans_in_file = read_spans_from_file(storage.filename)
+    assert len(spans_in_file) == 2
+    assert spans_in_file[0]["trace_id"] == "1"
+    assert spans_in_file[1]["trace_id"] == "2"
+
+
+@pytest.mark.parametrize("storage", [{"file_does_not_exist": True}], indirect=True)
+def test_log_rollover(storage: SpanStorage):
+    # mock start date
+    mock_start_date = datetime(2023, 10, 18, 23, 59)
+
+    with patch("datetime.datetime") as mock_datetime:
+        mock_datetime.now.return_value = mock_start_date
+
+        span1 = Span("1", "a", SpanType.BASE, "b", "op1")
+        storage.append_span(span1)
+        time.sleep(0.1)
+
+        # mock new day
+        mock_datetime.now.return_value = mock_start_date + timedelta(minutes=1)
+
+        span2 = Span("2", "c", SpanType.BASE, "d", "op2")
+        storage.append_span(span2)
+        time.sleep(0.1)
+
+    # origin filename need exists
+    assert os.path.exists(storage.filename)
+
+    # get roll over filename
+    dated_filename = os.path.join(
+        os.path.dirname(storage.filename),
+        f"{os.path.basename(storage.filename).split('.')[0]}_2023-10-18.jsonl",
+    )
+
+    assert os.path.exists(dated_filename)
+
+    # check origin filename just include the second span
+    spans_in_original_file = read_spans_from_file(storage.filename)
+    assert len(spans_in_original_file) == 1
+    assert spans_in_original_file[0]["trace_id"] == "2"
+
+    # check the roll over filename just include the first span
+    spans_in_dated_file = read_spans_from_file(dated_filename)
+    assert len(spans_in_dated_file) == 1
+    assert spans_in_dated_file[0]["trace_id"] == "1"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("storage_container", [{"batch_size": 5}], indirect=True)
+async def test_container_flush_policy(
+    storage_container: SpanStorageContainer, storage: FileSpanStorage
+):
+    storage_container.append_storage(storage)
+    span = Span("1", "a", SpanType.BASE, "b", "op1")
+
+    filename = storage.filename
+
+    for _ in range(storage_container.batch_size - 1):
+        storage_container.append_span(span)
+
+    spans_in_file = read_spans_from_file(filename)
+    assert len(spans_in_file) == 0
+
+    # Trigger batch write
+    storage_container.append_span(span)
+    await asyncio.sleep(0.1)
+
+    spans_in_file = read_spans_from_file(filename)
+    assert len(spans_in_file) == storage_container.batch_size

dbgpt/util/tracer/tests/test_tracer_impl.py

@@ -0,0 +1,103 @@
+import pytest
+from dbgpt.util.tracer import (
+    Span,
+    SpanStorageType,
+    SpanStorage,
+    DefaultTracer,
+    TracerManager,
+    Tracer,
+    MemorySpanStorage,
+)
+from dbgpt.component import SystemApp
+
+
+@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
+
+
+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 in tracer._get_current_storage().spans
+    assert child_span 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