chore: Fix tests/races (#1439)

Zerolog ~(Mac only???)~ testwriter issue: https://github.com/cloudquery/plugin-sdk/actions/runs/7350939011/job/20013680798?pr=1437
Flakey test on usage: https://github.com/cloudquery/plugin-sdk/actions/runs/7351363762/job/20014584046?pr=1437

Doesn't really 'fix' the zerolog cancellation testwriter case but circumvents it.

I've tried doing the `logger.With().Logger()` dance all over the schedulers (for every goroutine) to no effect.

Author: disq

premium/usage_test.go

@@ -7,6 +7,7 @@ import (
 	"math"
 	"net/http"
 	"net/http/httptest"
+	"sync"
 	"testing"
 	"time"
 
@@ -408,6 +409,8 @@ func createTestServerWithRemainingRows(t *testing.T, remainingRows int) *testSta
 			err := dec.Decode(&req)
 			require.NoError(t, err)
 
+			stage.mu.Lock()
+			defer stage.mu.Unlock()
 			stage.update = append(stage.update, req.Rows)
 
 			w.WriteHeader(http.StatusOK)
@@ -429,13 +432,18 @@ type testStage struct {
 
 	remainingRows int
 	update        []int
+	mu            sync.RWMutex
 }
 
 func (s *testStage) numberOfUpdates() int {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
 	return len(s.update)
 }
 
 func (s *testStage) sumOfUpdates() int {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
 	sum := 0
 	for _, val := range s.update {
 		sum += val
@@ -444,6 +452,8 @@ func (s *testStage) sumOfUpdates() int {
 }
 
 func (s *testStage) minExcludingClose() int {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
 	m := math.MaxInt
 	for i := 0; i < len(s.update); i++ {
 		if s.update[i] < m {

scheduler/scheduler_test.go

@@ -3,8 +3,6 @@ package scheduler
 import (
 	"context"
 	"fmt"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 	"testing"
 
 	"github.com/apache/arrow/go/v15/arrow"
@@ -13,6 +11,8 @@ import (
 	"github.com/cloudquery/plugin-sdk/v4/scalar"
 	"github.com/cloudquery/plugin-sdk/v4/schema"
 	"github.com/rs/zerolog"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 type testExecutionClient struct {
@@ -276,10 +276,15 @@ func TestScheduler_Cancellation(t *testing.T) {
 	}
 
 	for _, strategy := range AllStrategies {
+		strategy := strategy
 		for _, tc := range tests {
 			tc := tc
 			t.Run(fmt.Sprintf("%s_%s", tc.name, strategy.String()), func(t *testing.T) {
-				sc := NewScheduler(WithLogger(zerolog.New(zerolog.NewTestWriter(t))), WithStrategy(strategy))
+				logger := zerolog.New(zerolog.NewTestWriter(t))
+				if tc.cancel {
+					logger = zerolog.Nop() // FIXME without this, zerolog usage causes a race condition when tests are run with `-race -count=100`
+				}
+				sc := NewScheduler(WithLogger(logger), WithStrategy(strategy))
 
 				messages := make(chan message.SyncMessage)
 				ctx, cancel := context.WithCancel(context.Background())