scripts/telemetry_poc.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.

"""Telemetry POC script."""

import argparse
import time
from typing import List, Optional

from chromite.lib import commandline
from chromite.utils.telemetry import exporter
from chromite.utils.telemetry import trace


def get_parser() -> commandline.ArgumentParser:
    """Build the argument parser."""
    parser = commandline.ArgumentParser(description=__doc__)

    parser.add_argument(
        "-t",
        "--time",
        type=int,
        default=1,
        help="Seconds to sleep.",
    )

    return parser


def parse_arguments(argv: List) -> argparse.Namespace:
    """Parse and validate arguments."""
    parser = get_parser()
    opts = parser.parse_args(argv)

    opts.Freeze()
    return opts


def main(argv: Optional[List[str]]) -> Optional[int]:
    """Main."""
    opts = parse_arguments(argv)

    clearcut = exporter.ClearcutSpanExporter()
    tracer = trace.PocTracer()
    tracer.exporters.append(clearcut)
    with tracer:
        with tracer.start_as_current_span("test") as span:
            time.sleep(opts.time / 2)
            span.add_event(name="mid-sleep-event", attrs={"attr": "val"})
            time.sleep(opts.time / 2)
        span.set_status(code=0)