From e086fd3007508ada3ab1b011dc6230f175ca4ec3 Mon Sep 17 00:00:00 2001
From: Michael Gaffney <mike@gaffney.cc>
Date: Tue, 19 Nov 2024 08:49:57 -0500
Subject: [PATCH] test(batch): Fix data races in tests

Fixes: 9b984b0ab632f9003c80bf0b6637eb8f4ce55c23
---
 internal/scheduler/batch/batch_test.go | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/internal/scheduler/batch/batch_test.go b/internal/scheduler/batch/batch_test.go
index b82352b6a..8fc60f405 100644
--- a/internal/scheduler/batch/batch_test.go
+++ b/internal/scheduler/batch/batch_test.go
@@ -7,6 +7,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"sync"
 	"testing"
 	"time"
 
@@ -300,18 +301,25 @@ type recorder struct {
 	execBatchSize  int
 	storeBatchSize int
 	status         scheduler.JobStatus
+	mu             sync.Mutex
 }
 
 func (r *recorder) setup(c *Config) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
 	c.Store = r.Store
 }
 
 func (r *recorder) Store(ctx context.Context, batchSize int) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
 	r.storeBatchSize = batchSize
 	return nil
 }
 
 func (r *recorder) Exec(ctx context.Context, batchSize int) (int, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
 	r.execBatchSize = batchSize
 	return 0, nil
 }
@@ -320,6 +328,7 @@ type testRun struct {
 	ret func(context.Context, int, *Config) (int, error)
 	chk func(*testing.T, *recorder)
 	rec *recorder
+	mu  sync.Mutex
 }
 
 func (tr *testRun) validate(t *testing.T) {
@@ -329,6 +338,8 @@ func (tr *testRun) validate(t *testing.T) {
 }
 
 func (tr *testRun) recorder(cf *Config) *recorder {
+	tr.mu.Lock()
+	defer tr.mu.Unlock()
 	if tr.rec == nil {
 		tr.rec = &recorder{}
 		tr.rec.setup(cf)