improve GoDI test assertion

pkg/dynamicinstrumentation/diconfig/binary_inspection.go

@@ -27,7 +27,7 @@ func inspectGoBinaries(configEvent ditypes.DIProcs) error {
 	for i := range configEvent {
 		err = AnalyzeBinary(configEvent[i])
 		if err != nil {
-			log.Info("inspection of PID %d (path=%s) failed: %w", configEvent[i].PID, configEvent[i].BinaryPath, err)
+			log.Infof("inspection of PID %d (path=%s) failed: %s", configEvent[i].PID, configEvent[i].BinaryPath, err)
 		} else {
 			inspectedAtLeastOneBinary = true
 		}

pkg/dynamicinstrumentation/proctracker/proctracker.go

@@ -20,6 +20,8 @@ import (
 
 	"golang.org/x/sys/unix"
 
+	delve "github.com/go-delve/delve/pkg/goversion"
+
 	"github.com/DataDog/datadog-agent/pkg/dynamicinstrumentation/ditypes"
 	"github.com/DataDog/datadog-agent/pkg/network/go/bininspect"
 	"github.com/DataDog/datadog-agent/pkg/process/monitor"
@@ -28,7 +30,6 @@ import (
 	"github.com/DataDog/datadog-agent/pkg/util/kernel"
 	"github.com/DataDog/datadog-agent/pkg/util/log"
 	"github.com/DataDog/datadog-agent/pkg/util/safeelf"
-	delve "github.com/go-delve/delve/pkg/goversion"
 )
 
 type processTrackerCallback func(ditypes.DIProcs)
@@ -80,6 +81,7 @@ func (pt *ProcessTracker) Stop() {
 	for _, unsubscribe := range pt.unsubscribe {
 		unsubscribe()
 	}
+	pt.pm.Stop()
 }
 
 func (pt *ProcessTracker) handleProcessStart(pid uint32) {

pkg/dynamicinstrumentation/testutil/e2e_test.go

@@ -10,46 +10,38 @@ package testutil
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
+	"maps"
 	"math/rand"
 	"os"
 	"os/exec"
-	"reflect"
-	"sort"
+	"slices"
 	"strings"
 	"testing"
 	"text/template"
 	"time"
 
+	"github.com/cilium/ebpf"
+	"github.com/cilium/ebpf/features"
+	"github.com/cilium/ebpf/rlimit"
 	"github.com/kr/pretty"
+	"github.com/stretchr/testify/assert"
 
 	"github.com/DataDog/datadog-agent/pkg/dynamicinstrumentation"
+	"github.com/DataDog/datadog-agent/pkg/dynamicinstrumentation/diagnostics"
 	"github.com/DataDog/datadog-agent/pkg/dynamicinstrumentation/diconfig"
 	"github.com/DataDog/datadog-agent/pkg/dynamicinstrumentation/ditypes"
 	"github.com/DataDog/datadog-agent/pkg/util/testutil/flake"
-	"github.com/cilium/ebpf"
-	"github.com/cilium/ebpf/features"
-	"github.com/cilium/ebpf/rlimit"
 
 	"github.com/stretchr/testify/require"
 )
 
-type testResult struct {
-	testName          string
-	matches           []bool
-	expectation       ditypes.CapturedValueMap
-	unexpectedResults []ditypes.CapturedValueMap
-}
-
-var results = make(map[string]*testResult)
-
 func TestGoDI(t *testing.T) {
 	flake.Mark(t)
-	if err := rlimit.RemoveMemlock(); err != nil {
-		require.NoError(t, rlimit.RemoveMemlock())
-	}
-
+	require.NoError(t, rlimit.RemoveMemlock())
 	if features.HaveMapType(ebpf.RingBuf) != nil {
 		t.Skip("ringbuffers not supported on this kernel")
 	}
@@ -65,38 +57,52 @@ func TestGoDI(t *testing.T) {
 }
 
 func runTestCase(t *testing.T, function string, expectedCaptureValue CapturedValueMapWithOptions) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
 	serviceName := "go-di-sample-service-" + randomLabel()
 	sampleServicePath := BuildSampleService(t)
-	cmd := exec.Command(sampleServicePath)
+	cmd := exec.CommandContext(ctx, sampleServicePath)
+	cmd.WaitDelay = 10 * time.Millisecond
 	cmd.Env = []string{
 		"DD_DYNAMIC_INSTRUMENTATION_ENABLED=true",
 		fmt.Sprintf("DD_SERVICE=%s", serviceName),
 		"DD_DYNAMIC_INSTRUMENTATION_OFFLINE=true",
 	}
 
-	stdoutPipe, err1 := cmd.StdoutPipe()
-	require.NoError(t, err1)
+	stdoutPipe, err := cmd.StdoutPipe()
+	require.NoError(t, err)
 	go func() {
 		scanner := bufio.NewScanner(stdoutPipe)
 		for scanner.Scan() {
-			t.Log(scanner.Text())
+			t.Log("sample_service:", scanner.Text())
 		}
-		if err := scanner.Err(); err != nil {
-			t.Log(err)
+		if err := scanner.Err(); err != nil && !errors.Is(err, os.ErrClosed) {
+			t.Log("sample_service:", err)
 		}
 	}()
+	t.Cleanup(func() {
+		stdoutPipe.Close()
+	})
 
 	// send stderr to stdout pipe
 	cmd.Stderr = cmd.Stdout
 
 	require.NoError(t, cmd.Start())
 	t.Cleanup(func() {
-		cmd.Process.Kill()
+		t.Logf("stopping %d", cmd.Process.Pid)
+		_ = cmd.Process.Signal(os.Interrupt)
+		_ = cmd.Wait()
 	})
+	t.Logf("launched %s as %d", serviceName, cmd.Process.Pid)
 
 	eventOutputWriter := &eventOutputTestWriter{
-		t: t,
+		snapshots: make(chan ditypes.SnapshotUpload, 100),
 	}
+	t.Cleanup(func() { close(eventOutputWriter.snapshots) })
+
+	diagnostics.Diagnostics = diagnostics.NewDiagnosticManager()
+	t.Cleanup(diagnostics.StopGlobalDiagnostics)
 
 	opts := &dynamicinstrumentation.DIOptions{
 		RateLimitPerProbePerSecond: 0.0,
@@ -107,73 +113,53 @@ func runTestCase(t *testing.T, function string, expectedCaptureValue CapturedVal
 		},
 	}
 
-	var (
-		GoDI *dynamicinstrumentation.GoDI
-		err  error
-	)
-
-	GoDI, err = dynamicinstrumentation.RunDynamicInstrumentation(opts)
+	GoDI, err := dynamicinstrumentation.RunDynamicInstrumentation(opts)
 	require.NoError(t, err)
 	t.Cleanup(GoDI.Close)
 
 	cm, ok := GoDI.ConfigManager.(*diconfig.ReaderConfigManager)
-	if !ok {
-		t.Fatal("Config manager is of wrong type")
-	}
+	require.True(t, ok, "Config manager is of wrong type")
 
 	cfgTemplate, err := template.New("config_template").Parse(configTemplateText)
-	require.NoError(t, err)
-
-	b := []byte{}
-	buf := bytes.NewBuffer(b)
+	require.NoError(t, err, "template parse")
 
 	// Generate config for this function
-	functionWithoutPackagePrefix, _ := strings.CutPrefix(function, "github.com/DataDog/datadog-agent/pkg/dynamicinstrumentation/testutil/sample.")
-	t.Log("Instrumenting ", functionWithoutPackagePrefix)
-	result := &testResult{
-		testName:          function,
-		expectation:       expectedCaptureValue.CapturedValueMap,
-		matches:           []bool{},
-		unexpectedResults: []ditypes.CapturedValueMap{},
-	}
-	results[function] = result
+	functionWithoutPackagePrefix := strings.TrimPrefix(function, "github.com/DataDog/datadog-agent/pkg/dynamicinstrumentation/testutil/sample.")
+	t.Log("Instrumenting", functionWithoutPackagePrefix)
+	buf := &bytes.Buffer{}
 	err = cfgTemplate.Execute(buf, configDataType{
 		ServiceName:  serviceName,
 		FunctionName: functionWithoutPackagePrefix,
 		CaptureDepth: expectedCaptureValue.Options.CaptureDepth,
 	})
-	require.NoError(t, err)
-	eventOutputWriter.expectedResult = expectedCaptureValue.CapturedValueMap
+	require.NoError(t, err, "template execute")
 
 	// Read the configuration via the config manager
 	_, err = cm.ConfigWriter.Write(buf.Bytes())
-	time.Sleep(time.Second * 2)
-	doCapture = true
-	if err != nil {
-		t.Errorf("could not read new configuration: %s", err)
-	}
-	time.Sleep(time.Second * 2)
-	doCapture = false
+	require.NoError(t, err, "read new configuration")
 
-	// Check results
-	for _, ok := range result.matches {
+	var lastSnapshot ditypes.CapturedValueMap
+	assert.EventuallyWithT(t, func(t *assert.CollectT) {
+		snapshot, ok := <-eventOutputWriter.snapshots
 		if !ok {
-			t.Errorf("Failed test for: %s\nReceived event: %v\nExpected: %v",
-				result.testName,
-				pretty.Sprint(result.unexpectedResults),
-				pretty.Sprint(result.expectation))
-			break
+			return
 		}
+		actual := snapshot.Debugger.Captures.Entry.Arguments
+		scrubPointerValues(actual)
+		compareCapturedValues(t, "", expectedCaptureValue.CapturedValueMap, actual)
+		lastSnapshot = actual
+	}, 5*time.Second, 100*time.Millisecond)
+
+	if t.Failed() && lastSnapshot != nil {
+		t.Logf("Expected:\n%s", pretty.Sprint(expectedCaptureValue.CapturedValueMap))
+		t.Logf("Actual:\n%s", pretty.Sprint(lastSnapshot))
 	}
 }
 
 type eventOutputTestWriter struct {
-	t              *testing.T
-	expectedResult map[string]*ditypes.CapturedValue
+	snapshots chan ditypes.SnapshotUpload
 }
 
-var doCapture bool
-
 // compareCapturedValues compares two CapturedValueMap objects in a deterministic way.
 // This function is needed because the test results are stored in maps, which don't guarantee
 // a consistent iteration order.
@@ -183,80 +169,62 @@ var doCapture bool
 // 2. Sorting keys before comparison
 // 3. Recursively comparing nested fields
 // 4. Comparing all relevant fields (Type, NotCapturedReason, Value) in a deterministic order
-func compareCapturedValues(expected, actual ditypes.CapturedValueMap) bool {
-	if len(expected) != len(actual) {
-		return false
-	}
-
-	expectedKeys := make([]string, 0, len(expected))
-	for k := range expected {
-		expectedKeys = append(expectedKeys, k)
-	}
-	sort.Strings(expectedKeys)
-
-	actualKeys := make([]string, 0, len(actual))
-	for k := range actual {
-		actualKeys = append(actualKeys, k)
-	}
-	sort.Strings(actualKeys)
-
-	if !reflect.DeepEqual(expectedKeys, actualKeys) {
-		return false
-	}
+func compareCapturedValues(t assert.TestingT, path string, expected, actual ditypes.CapturedValueMap) {
+	expectedKeys := slices.Collect(maps.Keys(expected))
+	actualKeys := slices.Collect(maps.Keys(actual))
+	assert.ElementsMatch(t, expectedKeys, actualKeys, "map keys")
 
 	for _, k := range expectedKeys {
-		expectedVal := expected[k]
-		actualVal := actual[k]
-
-		if expectedVal.Type != actualVal.Type {
-			return false
-		}
-
-		if expectedVal.NotCapturedReason != actualVal.NotCapturedReason {
-			return false
+		expectedVal, eok := expected[k]
+		actualVal, aok := actual[k]
+		if !eok || !aok {
+			continue
 		}
 
-		if expectedVal.Value != nil && actualVal.Value != nil {
-			if *expectedVal.Value != *actualVal.Value {
-				return false
-			}
-		} else if expectedVal.Value != nil || actualVal.Value != nil {
-			return false
+		var prefix string
+		if path != "" {
+			prefix = fmt.Sprintf("%s.%s", path, k)
+		} else {
+			prefix = k
 		}
+		compareCapturedValue(t, prefix, expectedVal, actualVal)
+	}
+}
 
-		if !compareCapturedValues(expectedVal.Fields, actualVal.Fields) {
-			return false
+func compareCapturedValue(t assert.TestingT, path string, expected, actual *ditypes.CapturedValue) {
+	assert.Equal(t, expected.Type, actual.Type, "Path: %q\nField: Type", path)
+	assert.Equal(t, stringComparePtrValue(expected.Value), stringComparePtrValue(actual.Value), "Path: %q\nField: Value", path)
+	compareCapturedValues(t, path, expected.Fields, actual.Fields)
+	// Entries seems unused, so don't compare
+	assert.Len(t, actual.Elements, len(expected.Elements), "Path: %q\nField: Elements (length)", path)
+	for i := range expected.Elements {
+		if i >= len(actual.Elements) {
+			continue
 		}
+		compareCapturedValue(t, fmt.Sprintf("%s.Elements[%d]", path, i), &expected.Elements[i], &actual.Elements[i])
 	}
 
-	return true
+	assert.Equal(t, expected.NotCapturedReason, actual.NotCapturedReason, "Path: %q\nField: NotCapturedReason", path)
+	assert.Equal(t, expected.IsNull, actual.IsNull, "Path: %q\nField: IsNull", path)
+	assert.Equal(t, expected.Size, actual.Size, "Path: %q\nField: Size", path)
+	assert.Equal(t, expected.Truncated, actual.Truncated, "Path: %q\nField: Truncated", path)
 }
 
-func (e *eventOutputTestWriter) Write(p []byte) (n int, err error) {
-	if !doCapture {
-		return 0, nil
-	}
-	var snapshot ditypes.SnapshotUpload
-	if err := json.Unmarshal(p, &snapshot); err != nil {
-		e.t.Error("failed to unmarshal snapshot", err)
+func stringComparePtrValue(x *string) any {
+	if x == nil {
+		return nil
 	}
+	return *x
+}
 
-	funcName := snapshot.Debugger.ProbeInSnapshot.Method
-	actual := snapshot.Debugger.Captures.Entry.Arguments
-	scrubPointerValues(actual)
-	b, ok := results[funcName]
-	if !ok {
-		e.t.Errorf("received event from unexpected probe: %s", funcName)
-		return
-	}
-	if !compareCapturedValues(e.expectedResult, actual) {
-		b.matches = append(b.matches, false)
-		b.unexpectedResults = append(b.unexpectedResults, actual)
-		e.t.Error("received unexpected value")
-	} else {
-		b.matches = append(b.matches, true)
+func (e *eventOutputTestWriter) Write(p []byte) (n int, err error) {
+	var snapshot ditypes.SnapshotUpload
+	err = json.Unmarshal(p, &snapshot)
+	if err != nil {
+		return 0, err
 	}
 
+	e.snapshots <- snapshot
 	return len(p), nil
 }