chore(tests): fix data races in streaming API intg tests (#8059)

master/internal/api_experiment_intg_test.go

@@ -9,6 +9,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"sort"
+	"sync"
 	"testing"
 	"time"
 	"unsafe"
@@ -50,11 +51,15 @@ import (
 )
 
 type mockStream[T any] struct {
+	mu   sync.Mutex
 	ctx  context.Context
 	data []T
 }
 
 func (m *mockStream[T]) Send(resp T) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
 	m.data = append(m.data, resp)
 	return nil
 }
@@ -65,6 +70,16 @@ func (m *mockStream[T]) Context() context.Context      { return m.ctx }
 func (m *mockStream[T]) SendMsg(mes interface{}) error { return nil }
 func (m *mockStream[T]) RecvMsg(mes interface{}) error { return nil }
 
+func (m *mockStream[T]) getData() []T {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	out := make([]T, len(m.data))
+	copy(out, m.data)
+
+	return out
+}
+
 var (
 	authZExp   *mocks.ExperimentAuthZ
 	authzModel *mocks.ModelAuthZ

master/internal/api_trials_intg_test.go

@@ -377,7 +377,7 @@ func TestStreamTrainingMetrics(t *testing.T) {
 					return nil, err
 				}
 				var out []*trialv1.MetricsReport
-				for _, d := range res.data {
+				for _, d := range res.getData() {
 					out = append(out, d.Metrics...)
 				}
 				return out, nil
@@ -393,7 +393,7 @@ func TestStreamTrainingMetrics(t *testing.T) {
 					return nil, err
 				}
 				var out []*trialv1.MetricsReport
-				for _, d := range res.data {
+				for _, d := range res.getData() {
 					out = append(out, d.Metrics...)
 				}
 				return out, nil
@@ -503,7 +503,7 @@ func TestTrialsNonNumericMetrics(t *testing.T) {
 			resp := &mockStream[*apiv1.TrialsSampleResponse]{ctx: childCtx}
 			go func() {
 				for i := 0; i < 100; i++ {
-					if len(resp.data) > 0 {
+					if len(resp.getData()) > 0 {
 						cancel()
 					}
 					time.Sleep(50 * time.Millisecond)
@@ -519,9 +519,10 @@ func TestTrialsNonNumericMetrics(t *testing.T) {
 			}, resp)
 			require.NoError(t, err)
 
-			require.Greater(t, len(resp.data), 0)
-			require.Len(t, resp.data[0].Trials, 1)
-			require.Len(t, resp.data[0].Trials[0].Data, 1)
+			data := resp.getData()
+			require.Greater(t, len(data), 0)
+			require.Len(t, data[0].Trials, 1)
+			require.Len(t, data[0].Trials[0].Data, 1)
 			require.Equal(t, map[string]any{
 				metricName: expectedMetricsMap[metricName],
 			}, resp.data[0].Trials[0].Data[0].Values.AsMap())
@@ -871,9 +872,10 @@ func TestTrialLogsBackported(t *testing.T) {
 	}, stream)
 	require.NoError(t, err)
 
-	require.Equal(t, len(expected), len(stream.data))
+	actual := stream.getData()
+	require.Equal(t, len(expected), len(actual))
 	for i, expected := range expected {
-		require.Equal(t, expected.Log, *stream.data[i].Log)
+		require.Equal(t, expected.Log, *actual[i].Log)
 	}
 }
 
@@ -966,9 +968,10 @@ func TestTrialLogs(t *testing.T) {
 		t.Fatal("follow is following too long task logs")
 	}
 
-	require.Equal(t, len(expected), len(newStream.data))
+	actual := newStream.getData()
+	require.Equal(t, len(expected), len(actual))
 	for i, expected := range expected {
-		require.Equal(t, expected, *newStream.data[i].Log)
+		require.Equal(t, expected, *actual[i].Log)
 	}
 }
 
@@ -1020,7 +1023,7 @@ func TestTrialLogFields(t *testing.T) {
 	require.NoError(t, err)
 
 	actualContainerIDs := make(map[string]bool)
-	for _, s := range stream.data {
+	for _, s := range stream.getData() {
 		for _, containerID := range s.ContainerIds {
 			actualContainerIDs[containerID] = true
 		}
@@ -1071,7 +1074,7 @@ func TestTrialLogFields(t *testing.T) {
 	}
 
 	actualContainerIDs = make(map[string]bool)
-	for _, s := range newStream.data {
+	for _, s := range newStream.getData() {
 		for _, containerID := range s.ContainerIds {
 			actualContainerIDs[containerID] = true
 		}