lib/telemetry/trace/chromite_tracer_unittest.py

# Copyright 2023 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Test the trace implementation."""

import contextlib
from typing import Sequence

from chromite.third_party.opentelemetry import trace as trace_api
from chromite.third_party.opentelemetry.sdk import trace as trace_sdk
from chromite.third_party.opentelemetry.sdk.trace import export as export_sdk

from chromite.lib.telemetry.trace import chromite_tracer


class SpanExporterStub(export_sdk.SpanExporter):
    """Stub for SpanExporter."""

    def __init__(self) -> None:
        self._spans = []

    @property
    def spans(self):
        return self._spans

    def export(
        self, spans: Sequence[trace_sdk.ReadableSpan]
    ) -> export_sdk.SpanExportResult:
        self._spans.extend(spans)

        return export_sdk.SpanExportResult.SUCCESS


class ExceptionWithFailedPackages(Exception):
    """Exception with failed_packages."""

    def __init__(self, msg: str, pkgs: Sequence[str]) -> None:
        super().__init__(msg)

        self.failed_packages = pkgs


def test_chromite_span_to_capture_keyboard_interrupt_as_decorator() -> None:
    """Test chromite span can capture KeyboardInterrupt as decorator."""
    exporter = SpanExporterStub()
    provider = chromite_tracer.ChromiteTracerProvider(
        trace_sdk.TracerProvider()
    )
    provider.add_span_processor(
        export_sdk.BatchSpanProcessor(span_exporter=exporter)
    )
    tracer = provider.get_tracer(__name__)

    @tracer.start_as_current_span("test-span")
    def test_function() -> None:
        raise KeyboardInterrupt()

    with contextlib.suppress(KeyboardInterrupt):
        test_function()

    provider.shutdown()

    assert len(exporter.spans) == 1
    assert exporter.spans[0].events[0].name == "exception"
    assert (
        exporter.spans[0].events[0].attributes["exception.type"]
        == "KeyboardInterrupt"
    )
    assert exporter.spans[0].status.status_code == trace_api.StatusCode.OK


def test_chromite_span_to_capture_keyboard_interrupt_in_context() -> None:
    """Test chromite span can capture KeyboardInterrupt in context."""
    exporter = SpanExporterStub()
    provider = chromite_tracer.ChromiteTracerProvider(
        trace_sdk.TracerProvider()
    )
    provider.add_span_processor(
        export_sdk.BatchSpanProcessor(span_exporter=exporter)
    )
    tracer = provider.get_tracer(__name__)

    with contextlib.suppress(KeyboardInterrupt):
        with tracer.start_as_current_span("test-span"):
            raise KeyboardInterrupt()

    provider.shutdown()

    assert len(exporter.spans) == 1
    assert exporter.spans[0].events[0].name == "exception"
    assert (
        exporter.spans[0].events[0].attributes["exception.type"]
        == "KeyboardInterrupt"
    )
    assert exporter.spans[0].status.status_code == trace_api.StatusCode.OK


def test_chromite_span_to_capture_failed_packages_in_context() -> None:
    """Test chromite span can capture failed_packages in context."""
    exporter = SpanExporterStub()
    provider = chromite_tracer.ChromiteTracerProvider(
        trace_sdk.TracerProvider()
    )
    provider.add_span_processor(
        export_sdk.BatchSpanProcessor(span_exporter=exporter)
    )
    tracer = provider.get_tracer(__name__)

    with contextlib.suppress(ExceptionWithFailedPackages):
        with tracer.start_as_current_span("test-span"):
            raise ExceptionWithFailedPackages(
                msg="testing", pkgs=["dev-python/boto"]
            )

    provider.shutdown()

    assert len(exporter.spans) == 1
    assert exporter.spans[0].status.status_code == trace_api.StatusCode.ERROR
    assert (
        exporter.spans[0].status.description
        == "ExceptionWithFailedPackages: testing"
    )
    assert list(exporter.spans[0].events[0].attributes["failed_packages"]) == [
        "dev-python/boto"
    ]
    assert (
        exporter.spans[0].events[0].attributes["exception.type"]
        == "ExceptionWithFailedPackages"
    )


def test_chromite_span_to_capture_failed_packages_as_decorator() -> None:
    """Test chromite span can capture failed_packages as decorator."""
    exporter = SpanExporterStub()
    provider = chromite_tracer.ChromiteTracerProvider(
        trace_sdk.TracerProvider()
    )
    provider.add_span_processor(
        export_sdk.BatchSpanProcessor(span_exporter=exporter)
    )
    tracer = provider.get_tracer(__name__)

    @tracer.start_as_current_span("test-span")
    def test_function() -> None:
        raise ExceptionWithFailedPackages(
            msg="testing", pkgs=["dev-python/boto"]
        )

    with contextlib.suppress(ExceptionWithFailedPackages):
        test_function()

    provider.shutdown()

    assert len(exporter.spans) == 1
    assert exporter.spans[0].status.status_code == trace_api.StatusCode.ERROR
    assert (
        exporter.spans[0].status.description
        == "ExceptionWithFailedPackages: testing"
    )
    assert list(exporter.spans[0].events[0].attributes["failed_packages"]) == [
        "dev-python/boto"
    ]
    assert (
        exporter.spans[0].events[0].attributes["exception.type"]
        == "ExceptionWithFailedPackages"
    )