scouter_tracing/exporter/
testing.rs

1use crate::exporter::ExporterType;
2use crate::exporter::SpanExporterBuilder;
3use crate::exporter::TraceError;
4use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
5/// Implementation for testing exporter used in unit testsuse opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
6use opentelemetry_proto::transform::common::tonic::ResourceAttributesWithSchema;
7use opentelemetry_proto::transform::trace::tonic::group_spans_by_resource_and_scope;
8use opentelemetry_sdk::{
9    error::OTelSdkResult,
10    trace::{SpanData, SpanExporter},
11};
12use pyo3::prelude::*;
13use scouter_types::{TraceBaggageRecord, TraceRecord, TraceServerRecord, TraceSpanRecord};
14use std::sync::{Arc, RwLock};
15
16#[derive(Debug)]
17pub struct TestRecords {
18    pub traces: Vec<TraceRecord>,
19    pub spans: Vec<TraceSpanRecord>,
20    pub baggage: Vec<TraceBaggageRecord>,
21}
22
23#[derive(Debug, Clone)]
24#[pyclass]
25pub struct TestSpanExporter {
26    records: Arc<RwLock<TestRecords>>,
27    batch_export: bool,
28}
29
30#[pymethods]
31impl TestSpanExporter {
32    #[new]
33    #[pyo3(signature = (batch_export=true))]
34    pub fn new(batch_export: bool) -> Self {
35        TestSpanExporter {
36            records: Arc::new(RwLock::new(TestRecords {
37                traces: Vec::new(),
38                spans: Vec::new(),
39                baggage: Vec::new(),
40            })),
41            batch_export,
42        }
43    }
44
45    #[getter]
46    pub fn traces(&self) -> Vec<TraceRecord> {
47        self.records.read().unwrap().traces.clone()
48    }
49
50    #[getter]
51    pub fn spans(&self) -> Vec<TraceSpanRecord> {
52        self.records.read().unwrap().spans.clone()
53    }
54
55    #[getter]
56    pub fn baggage(&self) -> Vec<TraceBaggageRecord> {
57        self.records.read().unwrap().baggage.clone()
58    }
59
60    pub fn clear(&self) {
61        let mut records = self.records.write().unwrap();
62        records.traces.clear();
63        records.spans.clear();
64        records.baggage.clear();
65    }
66}
67
68impl Default for TestSpanExporter {
69    fn default() -> Self {
70        Self::new(true)
71    }
72}
73
74impl SpanExporterBuilder for TestSpanExporter {
75    type Exporter = OtelTestSpanExporter;
76
77    fn export_type(&self) -> ExporterType {
78        ExporterType::Testing
79    }
80
81    fn sample_ratio(&self) -> Option<f64> {
82        Some(1.0)
83    }
84
85    fn batch_export(&self) -> bool {
86        self.batch_export
87    }
88
89    fn build_exporter(&self) -> Result<Self::Exporter, TraceError> {
90        Ok(OtelTestSpanExporter::new(self.records.clone()))
91    }
92}
93
94#[derive(Debug)]
95pub struct OtelTestSpanExporter {
96    records: Arc<RwLock<TestRecords>>,
97}
98
99impl OtelTestSpanExporter {
100    pub fn new(records: Arc<RwLock<TestRecords>>) -> Self {
101        OtelTestSpanExporter { records }
102    }
103}
104
105impl SpanExporter for OtelTestSpanExporter {
106    async fn export(&self, batch: Vec<SpanData>) -> OTelSdkResult {
107        // Here you would implement the logic to export spans to Scouter
108        let resource_spans =
109            group_spans_by_resource_and_scope(batch, &ResourceAttributesWithSchema::default());
110
111        let req = ExportTraceServiceRequest { resource_spans };
112
113        let record = TraceServerRecord {
114            request: req,
115            space: "test_space".to_string(),
116            name: "test_name".to_string(),
117            version: "test_version".to_string(),
118        };
119
120        let (traces, spans, baggage) = record
121            .to_records()
122            .map_err(|e| opentelemetry_sdk::error::OTelSdkError::InternalFailure(e.to_string()))?;
123
124        let mut records = self.records.write().unwrap();
125        records.traces.extend(traces);
126        records.spans.extend(spans);
127        records.baggage.extend(baggage);
128
129        Ok(())
130    }
131
132    fn shutdown(&mut self) -> OTelSdkResult {
133        Ok(())
134    }
135}
scouter_tracing/exporter/testing.rs

scouter_tracing/exporter/
testing.rs
