diff --git a/pkg/controlbridge/helpers_test.go b/pkg/controlbridge/helpers_test.go index ea4f9a524..1a26912a8 100644 --- a/pkg/controlbridge/helpers_test.go +++ b/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() diff --git a/pkg/sqlreplay/store/dir_watcher_test.go b/pkg/sqlreplay/store/dir_watcher_test.go index ee66a9199..c77007af3 100644 --- a/pkg/sqlreplay/store/dir_watcher_test.go +++ b/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()