controlbridge, sqlreplay: fix Go lint and flaky TestS3DirWatcher on rust

pkg/controlbridge/helpers_test.go

@@ -3,7 +3,7 @@
 
 // Shared test doubles. These outlived the legacy RouterAdapter deleted at
 // #223 Phase 2: the residual route-owner handler, the snapshot publisher and
-// the metering tests all still need a control peer and a handshake handler.
+// the metering tests all still need a control peer.
 
 package controlbridge
 
@@ -13,12 +13,8 @@ import (
 	"sync"
 	"testing"
 
-	"github.com/go-mysql-org/go-mysql/mysql"
-	"github.com/pingcap/tiproxy/pkg/balance/router"
 	controlpb "github.com/pingcap/tiproxy/pkg/controlbridge/pb"
 	"github.com/pingcap/tiproxy/pkg/controlbridge/transport"
-	"github.com/pingcap/tiproxy/pkg/proxy/backend"
-	pnet "github.com/pingcap/tiproxy/pkg/proxy/net"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/proto"
 )
@@ -68,82 +64,6 @@ func (peer *fakeSender) AllocateRequestID() (uint64, error) {
 	return peer.nextID, nil
 }
 
-type recordingHandler struct {
-	mu                   sync.Mutex
-	rt                   router.Router
-	response             *pnet.HandshakeResp
-	retryHandshakeErrors int
-	handshakeCalls       int
-	handshakeErr         error
-	trafficCalls         int
-	closeCalls           int
-	lastClientIn         uint64
-	lastClientOut        uint64
-}
-
-func (handler *recordingHandler) HandleHandshakeResp(_ backend.ConnContext, response *pnet.HandshakeResp) error {
-	handler.mu.Lock()
-	defer handler.mu.Unlock()
-	handler.response = response
-	return nil
-}
-
-func (handler *recordingHandler) HandleHandshakeErr(_ backend.ConnContext, _ *mysql.MyError) bool {
-	handler.mu.Lock()
-	defer handler.mu.Unlock()
-	if handler.retryHandshakeErrors > 0 {
-		handler.retryHandshakeErrors--
-		return true
-	}
-	return false
-}
-
-func (handler *recordingHandler) GetRouter(_ backend.ConnContext, _ *pnet.HandshakeResp) (router.Router, error) {
-	if handler.rt == nil {
-		return nil, errors.New("no router")
-	}
-	return handler.rt, nil
-}
-
-func (handler *recordingHandler) OnHandshake(_ backend.ConnContext, _ string, err error, _ backend.ErrorSource) {
-	handler.mu.Lock()
-	defer handler.mu.Unlock()
-	handler.handshakeCalls++
-	handler.handshakeErr = err
-}
-
-func (handler *recordingHandler) OnConnClose(_ backend.ConnContext, _ backend.ErrorSource) error {
-	handler.mu.Lock()
-	defer handler.mu.Unlock()
-	handler.closeCalls++
-	return nil
-}
-
-func (handler *recordingHandler) OnTraffic(ctx backend.ConnContext) {
-	handler.mu.Lock()
-	defer handler.mu.Unlock()
-	handler.trafficCalls++
-	handler.lastClientIn = ctx.ClientInBytes()
-	handler.lastClientOut = ctx.ClientOutBytes()
-}
-
-func (handler *recordingHandler) GetCapability() pnet.Capability {
-	return backend.SupportedServerCapabilities
-}
-func (handler *recordingHandler) GetServerVersion() string { return pnet.ServerVersion }
-
-func (handler *recordingHandler) handshakeCount() int {
-	handler.mu.Lock()
-	defer handler.mu.Unlock()
-	return handler.handshakeCalls
-}
-
-func (handler *recordingHandler) closeCount() int {
-	handler.mu.Lock()
-	defer handler.mu.Unlock()
-	return handler.closeCalls
-}
-
 func lastEnvelope(t *testing.T, peer *fakeSender) *controlpb.ControlEnvelope {
 	t.Helper()
 	peer.mu.Lock()

pkg/sqlreplay/store/dir_watcher_test.go

@@ -121,7 +121,7 @@ func TestS3DirWatcher(t *testing.T) {
 
 	dirWatcherPollInterval = time.Millisecond * 10
 
-	s3api.EXPECT().ListObjectsV2(gomock.Any(), gomock.Any()).MaxTimes(4).DoAndReturn(
+	s3api.EXPECT().ListObjectsV2(gomock.Any(), gomock.Any()).AnyTimes().DoAndReturn(
 		func(ctx context.Context, req *s3v2.ListObjectsV2Input, _ ...func(*s3v2.Options)) (*s3v2.ListObjectsV2Output, error) {
 			var retFiles []types.Object
 			files := currentFiles.Load()
@@ -176,8 +176,11 @@ func TestS3DirWatcher(t *testing.T) {
 	})
 
 	t.Run("Watch function should have the correct result", func(t *testing.T) {
+		var mu sync.Mutex
 		files := make(map[string]struct{})
 		w := NewDirWatcher(logger, "dir_watcher_test_", func(filename string) error {
+			mu.Lock()
+			defer mu.Unlock()
 			files[filename] = struct{}{}
 			return nil
 		}, s)
@@ -191,10 +194,12 @@ func TestS3DirWatcher(t *testing.T) {
 		}()
 
 		require.Eventually(t, func() bool {
+			mu.Lock()
+			defer mu.Unlock()
 			_, ok1 := files[newDir]
 			_, ok2 := files[newDir2]
 			return ok1 && ok2 && len(files) == 2
-		}, dirWatcherPollInterval*3, time.Millisecond)
+		}, 5*time.Second, time.Millisecond)
 
 		oldFiles := currentFiles.Load()
 		newFiles := slices.Clone(*oldFiles)
@@ -203,9 +208,11 @@ func TestS3DirWatcher(t *testing.T) {
 		currentFiles.Store(&newFiles)
 
 		require.Eventually(t, func() bool {
+			mu.Lock()
+			defer mu.Unlock()
 			_, ok := files[newDir3]
 			return ok && len(files) == 3
-		}, dirWatcherPollInterval*3, time.Millisecond)
+		}, 5*time.Second, time.Millisecond)
 
 		cancel()
 		wg.Wait()