From d758e25ffa6b04347928076f23302c753bd0c20b Mon Sep 17 00:00:00 2001
From: Tony Redondo <tony.redondo@datadoghq.com>
Date: Wed, 11 Sep 2024 15:36:33 +0200
Subject: [PATCH] internal/civisibility/integrations/gotesting: fixes an issue
 where we were caching the instrumented function from an original one, caching
 also the closure data of the function, now we don't catch the instrumented
 function. Also, added some tests to handle this scenario.

---
 .../integrations/gotesting/instrumentation.go | 30 ++++---------------
 .../integrations/gotesting/testing.go         |  9 +-----
 .../integrations/gotesting/testing_test.go    | 21 ++++++++++---
 3 files changed, 24 insertions(+), 36 deletions(-)

diff --git a/internal/civisibility/integrations/gotesting/instrumentation.go b/internal/civisibility/integrations/gotesting/instrumentation.go
index 6077c34e10..0098505219 100644
--- a/internal/civisibility/integrations/gotesting/instrumentation.go
+++ b/internal/civisibility/integrations/gotesting/instrumentation.go
@@ -27,9 +27,7 @@ import (
 
 type (
 	instrumentationMetadata struct {
-		IsInternal       bool
-		OriginalTest     *func(*testing.T)
-		InstrumentedTest *func(*testing.T)
+		IsInternal bool
 	}
 
 	ddTestItem struct {
@@ -57,11 +55,9 @@ var (
 func getInstrumentationMetadata(fn *runtime.Func) *instrumentationMetadata {
 	instrumentationMapMutex.RLock()
 	defer instrumentationMapMutex.RUnlock()
-
 	if v, ok := instrumentationMap[fn]; ok {
 		return v
 	}
-
 	return nil
 }
 
@@ -76,11 +72,9 @@ func setInstrumentationMetadata(fn *runtime.Func, metadata *instrumentationMetad
 func getCiVisibilityTest(tb testing.TB) *ddTestItem {
 	ciVisibilityTestsMutex.RLock()
 	defer ciVisibilityTestsMutex.RUnlock()
-
 	if v, ok := ciVisibilityTests[reflect.ValueOf(tb).UnsafePointer()]; ok {
 		return v
 	}
-
 	return nil
 }
 
@@ -131,16 +125,12 @@ func instrumentTestingTFunc(f func(*testing.T)) func(*testing.T) {
 
 	// Avoid instrumenting twice
 	metadata := getInstrumentationMetadata(originalFunc)
-	if metadata != nil {
+	if metadata != nil && metadata.IsInternal {
 		// If is an internal test, we don't instrument because f is already the instrumented func by executeInternalTest
-		if metadata.IsInternal {
-			return f
-		}
-
-		return *metadata.InstrumentedTest
+		return f
 	}
 
-	instrumentedFunc := func(t *testing.T) {
+	instrumentedFn := func(t *testing.T) {
 		// Initialize module counters if not already present.
 		if _, ok := modulesCounters[moduleName]; !ok {
 			var v int32
@@ -190,16 +180,8 @@ func instrumentTestingTFunc(f func(*testing.T)) func(*testing.T) {
 		// Execute the original test function.
 		f(t)
 	}
-
-	metadata = &instrumentationMetadata{
-		IsInternal:       false,
-		OriginalTest:     &f,
-		InstrumentedTest: &instrumentedFunc,
-	}
-
-	setInstrumentationMetadata(originalFunc, metadata)
-	setInstrumentationMetadata(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()), metadata)
-	return instrumentedFunc
+	setInstrumentationMetadata(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFn)).Pointer()), &instrumentationMetadata{IsInternal: true})
+	return instrumentedFn
 }
 
 // instrumentSetErrorInfo helper function to set an error in the `testing.T or testing.B` CI Visibility span
diff --git a/internal/civisibility/integrations/gotesting/testing.go b/internal/civisibility/integrations/gotesting/testing.go
index 1722248dd9..b7ac9c381b 100644
--- a/internal/civisibility/integrations/gotesting/testing.go
+++ b/internal/civisibility/integrations/gotesting/testing.go
@@ -166,14 +166,7 @@ func (ddm *M) executeInternalTest(testInfo *testingTInfo) func(*testing.T) {
 		testInfo.originalFunc(t)
 	}
 
-	metadata := &instrumentationMetadata{
-		IsInternal:       true,
-		OriginalTest:     &testInfo.originalFunc,
-		InstrumentedTest: &instrumentedFunc,
-	}
-
-	setInstrumentationMetadata(originalFunc, metadata)
-	setInstrumentationMetadata(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()), metadata)
+	setInstrumentationMetadata(originalFunc, &instrumentationMetadata{IsInternal: true})
 	return instrumentedFunc
 }
 
diff --git a/internal/civisibility/integrations/gotesting/testing_test.go b/internal/civisibility/integrations/gotesting/testing_test.go
index e45e62d15d..4345e9e822 100644
--- a/internal/civisibility/integrations/gotesting/testing_test.go
+++ b/internal/civisibility/integrations/gotesting/testing_test.go
@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"slices"
 	"strconv"
 	"testing"
 
@@ -35,7 +36,10 @@ func TestMain(m *testing.M) {
 	// or use a helper method gotesting.RunM(m)
 
 	// os.Exit((*M)(m).Run())
-	_ = RunM(m)
+	exit := RunM(m)
+	if exit != 0 {
+		os.Exit(exit)
+	}
 
 	finishedSpans := mTracer.FinishedSpans()
 	// 1 session span
@@ -107,19 +111,28 @@ func Test_Foo(gt *testing.T) {
 	assertTest(gt)
 	t := (*T)(gt)
 	var tests = []struct {
+		index byte
 		name  string
 		input string
 		want  string
 	}{
-		{"yellow should return color", "yellow", "color"},
-		{"banana should return fruit", "banana", "fruit"},
-		{"duck should return animal", "duck", "animal"},
+		{1, "yellow should return color", "yellow", "color"},
+		{2, "banana should return fruit", "banana", "fruit"},
+		{3, "duck should return animal", "duck", "animal"},
 	}
+	buf := []byte{}
 	for _, test := range tests {
+		test := test
 		t.Run(test.name, func(t *testing.T) {
 			t.Log(test.name)
+			buf = append(buf, test.index)
 		})
 	}
+
+	expected := []byte{1, 2, 3}
+	if !slices.Equal(buf, expected) {
+		t.Error("error in subtests closure")
+	}
 }
 
 // TestWithExternalCalls demonstrates testing with external HTTP calls.