From 575bac12072fa9cca4eb5552b67be0f2d0abf119 Mon Sep 17 00:00:00 2001
From: Chris Smith <chrisdsmith@google.com>
Date: Fri, 3 Nov 2023 16:38:57 -0600
Subject: [PATCH] add internal/testutil/trace_otel.go

---
 internal/testutil/trace_otel.go | 54 +++++++++++++++++++++++++++++++++
 internal/trace/trace_test.go    | 15 +++------
 2 files changed, 58 insertions(+), 11 deletions(-)
 create mode 100644 internal/testutil/trace_otel.go

diff --git a/internal/testutil/trace_otel.go b/internal/testutil/trace_otel.go
new file mode 100644
index 000000000000..dd132ecb9953
--- /dev/null
+++ b/internal/testutil/trace_otel.go
@@ -0,0 +1,54 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package testutil
+
+import (
+	"context"
+
+	"go.opentelemetry.io/otel"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
+)
+
+// TestExporter is a test utility exporter. It should be created with NewOpenTelemetryTestExporter.
+type OpenTelemetryTestExporter struct {
+	exporter *tracetest.InMemoryExporter
+	tp       *sdktrace.TracerProvider
+}
+
+// NewOpenTelemetryTestExporter creates a OpenTelemetryTestExporter with
+// underlying InMemoryExporter and TracerProvider from OpenTelemetry.
+func NewOpenTelemetryTestExporter() *OpenTelemetryTestExporter {
+	exporter := tracetest.NewInMemoryExporter()
+	tp := sdktrace.NewTracerProvider(
+		sdktrace.WithSyncer(exporter),
+		sdktrace.WithSampler(sdktrace.AlwaysSample()),
+	)
+	otel.SetTracerProvider(tp)
+	return &OpenTelemetryTestExporter{
+		exporter: exporter,
+		tp:       tp,
+	}
+}
+
+// Spans returns the current in-memory stored spans.
+func (te *OpenTelemetryTestExporter) Spans() tracetest.SpanStubs {
+	return te.exporter.GetSpans()
+}
+
+// Unregister shuts down the underlying OpenTelemetry TracerProvider.
+func (te *OpenTelemetryTestExporter) Unregister(ctx context.Context) {
+	te.tp.Shutdown(ctx)
+}
diff --git a/internal/trace/trace_test.go b/internal/trace/trace_test.go
index 161f990d6028..a5caa98fcba0 100644
--- a/internal/trace/trace_test.go
+++ b/internal/trace/trace_test.go
@@ -25,11 +25,9 @@ import (
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/googleapis/gax-go/v2/apierror"
 	octrace "go.opencensus.io/trace"
-	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 	otcodes "go.opentelemetry.io/otel/codes"
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
-	"go.opentelemetry.io/otel/sdk/trace/tracetest"
 	"google.golang.org/api/googleapi"
 	"google.golang.org/genproto/googleapis/rpc/code"
 	"google.golang.org/grpc/codes"
@@ -103,15 +101,10 @@ func TestStartSpan_OpenTelemetry(t *testing.T) {
 	}()
 	TelemetryPlatformTracing = "opentelemetry"
 
-	exporter := tracetest.NewInMemoryExporter()
-	tp := sdktrace.NewTracerProvider(
-		sdktrace.WithSyncer(exporter),
-		sdktrace.WithSampler(sdktrace.AlwaysSample()),
-	)
-	defer tp.Shutdown(context.Background())
-	otel.SetTracerProvider(tp)
-
 	ctx := context.Background()
+	te := testutil.NewOpenTelemetryTestExporter()
+	defer te.Unregister(ctx)
+
 	ctx = StartSpan(ctx, "test-span")
 
 	TracePrintf(ctx, attrMap(), "Add my annotations")
@@ -125,7 +118,7 @@ func TestStartSpan_OpenTelemetry(t *testing.T) {
 	if !IsOpenTelemetryTracingEnabled() {
 		t.Errorf("got false, want true")
 	}
-	spans := exporter.GetSpans()
+	spans := te.Spans()
 	if len(spans) != 1 {
 		t.Fatalf("got %d, want 1", len(spans))
 	}