fix tests

Author: ndyakov

internal/pool/hooks_test.go

@@ -17,11 +17,12 @@ type TestHook struct {
 	PutError       error
 	ShouldPool     bool
 	ShouldRemove   bool
+	ShouldAccept   bool
 }
 
-func (th *TestHook) OnGet(ctx context.Context, conn *Conn, isNewConn bool) error {
+func (th *TestHook) OnGet(ctx context.Context, conn *Conn, isNewConn bool) (bool, error) {
 	th.OnGetCalled++
-	return th.GetError
+	return th.ShouldAccept, th.GetError
 }
 
 func (th *TestHook) OnPut(ctx context.Context, conn *Conn) (shouldPool bool, shouldRemove bool, err error) {
@@ -42,8 +43,8 @@ func TestPoolHookManager(t *testing.T) {
 	}
 
 	// Add hooks
-	hook1 := &TestHook{ShouldPool: true}
-	hook2 := &TestHook{ShouldPool: true}
+	hook1 := &TestHook{ShouldPool: true, ShouldAccept: true}
+	hook2 := &TestHook{ShouldPool: true, ShouldAccept: true}
 
 	manager.AddHook(hook1)
 	manager.AddHook(hook2)
@@ -107,11 +108,12 @@ func TestHookErrorHandling(t *testing.T) {
 
 	// Hook that returns error on Get
 	errorHook := &TestHook{
-		GetError:   errors.New("test error"),
-		ShouldPool: true,
+		GetError:     errors.New("test error"),
+		ShouldPool:   true,
+		ShouldAccept: true,
 	}
 
-	normalHook := &TestHook{ShouldPool: true}
+	normalHook := &TestHook{ShouldPool: true, ShouldAccept: true}
 
 	manager.AddHook(errorHook)
 	manager.AddHook(normalHook)
@@ -145,9 +147,10 @@ func TestHookShouldRemove(t *testing.T) {
 	removeHook := &TestHook{
 		ShouldPool:   false,
 		ShouldRemove: true,
+		ShouldAccept: true,
 	}
 
-	normalHook := &TestHook{ShouldPool: true}
+	normalHook := &TestHook{ShouldPool: true, ShouldAccept: true}
 
 	manager.AddHook(removeHook)
 	manager.AddHook(normalHook)
@@ -181,7 +184,7 @@ func TestHookShouldRemove(t *testing.T) {
 func TestPoolWithHooks(t *testing.T) {
 	// Create a pool with hooks
 	hookManager := NewPoolHookManager()
-	testHook := &TestHook{ShouldPool: true}
+	testHook := &TestHook{ShouldPool: true, ShouldAccept: true}
 	hookManager.AddHook(testHook)
 
 	opt := &Options{
@@ -208,7 +211,7 @@ func TestPoolWithHooks(t *testing.T) {
 	}
 
 	// Test adding hook to pool
-	additionalHook := &TestHook{ShouldPool: true}
+	additionalHook := &TestHook{ShouldPool: true, ShouldAccept: true}
 	pool.AddPoolHook(additionalHook)
 
 	if pool.hookManager.GetHookCount() != 2 {

redis_test.go

@@ -854,6 +854,7 @@ var _ = Describe("Credentials Provider Priority", func() {
 				credentials: initialCreds,
 				updates:     updatesChan,
 			},
+			PoolSize: 1, // Force single connection to ensure reauth is tested
 		}
 
 		client = redis.NewClient(opt)
@@ -865,11 +866,16 @@ var _ = Describe("Credentials Provider Priority", func() {
 
 		// Update credentials
 		opt.StreamingCredentialsProvider.(*mockStreamingProvider).updates <- updatedCreds
-		// wrongpass
-		time.Sleep(10 * time.Millisecond)
-		Expect(client.Ping(context.Background()).Err()).To(HaveOccurred())
-		time.Sleep(10 * time.Millisecond)
-		Expect(recorder.Contains("AUTH updated_user")).To(BeTrue())
+
+		// Wait for reauth to complete and verify updated credentials are used
+		// We need to keep trying Ping until we see the updated AUTH command
+		// because the reauth happens asynchronously
+		Eventually(func() bool {
+			// wrongpass
+			_ = client.Ping(context.Background()).Err()
+			return recorder.Contains("AUTH updated_user")
+		}, "1s", "50ms").Should(BeTrue())
+
 		close(updatesChan)
 	})
 })