should be the last test failing

Signed-off-by: Eliott Bouhana <[email protected]>

Author: eliottness

contrib/internal/telemetrytest/telemetry_test.go

@@ -12,6 +12,8 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
+
 	"gopkg.in/DataDog/dd-trace-go.v1/contrib/gorilla/mux"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/telemetry"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/telemetry/telemetrytest"
@@ -23,14 +25,12 @@ import (
 // sends the correct data to the telemetry client.
 func TestIntegrationInfo(t *testing.T) {
 	// mux.NewRouter() uses the net/http and gorilla/mux integration
-	client := new(telemetrytest.MockClient)
-	client.On("AppStart").Return()
-	client.On("MarkIntegrationAsLoaded", telemetry.Integration{Name: "net/http", Version: "", Error: ""}).Return()
-	client.On("MarkIntegrationAsLoaded", telemetry.Integration{Name: "gorilla/mux", Version: "", Error: ""}).Return()
+	client := new(telemetrytest.RecordClient)
 	telemetry.StartApp(client)
 	_ = mux.NewRouter()
 
-	client.AssertExpectations(t)
+	assert.Contains(t, client.Integrations, telemetry.Integration{Name: "net/http", Version: "", Error: ""})
+	assert.Contains(t, client.Integrations, telemetry.Integration{Name: "gorilla/mux", Version: "", Error: ""})
 }
 
 type contribPkg struct {

ddtrace/opentelemetry/telemetry_test.go

@@ -74,9 +74,7 @@ func TestTelemetry(t *testing.T) {
 				t.Setenv(k, v)
 			}
 			telemetryClient := new(telemetrytest.RecordClient)
-			original := telemetry.GlobalClient()
-			telemetry.SwapClient(telemetryClient)
-			defer telemetry.SwapClient(original)
+			defer telemetry.MockClient(telemetryClient)()
 
 			p := NewTracerProvider()
 			p.Tracer("")

ddtrace/opentelemetry/tracer_test.go

@@ -196,9 +196,7 @@ func TestShutdownOnce(t *testing.T) {
 
 func TestSpanTelemetry(t *testing.T) {
 	telemetryClient := new(telemetrytest.RecordClient)
-	original := telemetry.GlobalClient()
-	telemetry.SwapClient(telemetryClient)
-	defer telemetry.SwapClient(original)
+	defer telemetry.MockClient(telemetryClient)()
 	tp := NewTracerProvider()
 	otel.SetTracerProvider(tp)
 	tr := otel.Tracer("")

ddtrace/opentracer/tracer_test.go

@@ -116,9 +116,7 @@ func TestExtractError(t *testing.T) {
 
 func TestSpanTelemetry(t *testing.T) {
 	telemetryClient := new(telemetrytest.RecordClient)
-	original := telemetry.GlobalClient()
-	telemetry.SwapClient(telemetryClient)
-	defer telemetry.SwapClient(original)
+	defer telemetry.MockClient(telemetryClient)()
 	opentracing.SetGlobalTracer(New())
 	_ = opentracing.StartSpan("opentracing.span")
 	assert.NotZero(t, telemetryClient.Count(telemetry.NamespaceTracers, "spans_created", telemetryTags).Get())

ddtrace/tracer/otel_dd_mappings_test.go

@@ -31,9 +31,7 @@ func TestAssessSource(t *testing.T) {
 	})
 	t.Run("both", func(t *testing.T) {
 		telemetryClient := new(telemetrytest.RecordClient)
-		original := telemetry.GlobalClient()
-		telemetry.SwapClient(telemetryClient)
-		defer telemetry.SwapClient(original)
+		defer telemetry.MockClient(telemetryClient)()
 		// DD_SERVICE prevails
 		t.Setenv("DD_SERVICE", "abc")
 		t.Setenv("OTEL_SERVICE_NAME", "123")

profiler/telemetry_test.go

@@ -15,19 +15,11 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func mockGlobalClient(client telemetry.Client) func() {
-	orig := telemetry.GlobalClient()
-	telemetry.SwapClient(client)
-	return func() {
-		telemetry.SwapClient(orig)
-	}
-}
-
 // Test that the profiler sends the correct telemetry information
 func TestTelemetryEnabled(t *testing.T) {
 	t.Run("tracer start, profiler start", func(t *testing.T) {
 		telemetryClient := new(telemetrytest.RecordClient)
-		defer mockGlobalClient(telemetryClient)()
+		defer telemetry.MockClient(telemetryClient)()
 
 		tracer.Start()
 		defer tracer.Stop()
@@ -44,7 +36,7 @@ func TestTelemetryEnabled(t *testing.T) {
 	})
 	t.Run("only profiler start", func(t *testing.T) {
 		telemetryClient := new(telemetrytest.RecordClient)
-		defer mockGlobalClient(telemetryClient)()
+		defer telemetry.MockClient(telemetryClient)()
 		Start(
 			WithProfileTypes(
 				HeapProfile,