fix

Author: ndyakov

async_handoff_integration_test.go

@@ -272,16 +272,28 @@ func TestEventDrivenHandoffIntegration(t *testing.T) {
 			t.Fatalf("Failed to get connection: %v", err)
 		}
 
-		conn.MarkForHandoff("new-endpoint:6379", 12345)
+		if err := conn.MarkForHandoff("new-endpoint:6379", 12345); err != nil {
+			t.Fatalf("Failed to mark connection for handoff: %v", err)
+		}
+
+		// Set a mock initialization function
+		conn.SetInitConnFunc(func(ctx context.Context, cn *pool.Conn) error {
+			return nil
+		})
+
 		testPool.Put(ctx, conn)
 
 		// Verify handoff was queued
 		if !processor.IsHandoffPending(conn) {
 			t.Error("Handoff should be queued in pending map")
 		}
 
+		// Give the handoff a moment to start processing
+		time.Sleep(50 * time.Millisecond)
+
 		// Shutdown processor gracefully
-		shutdownCtx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+		// Use a longer timeout to account for slow dialer (100ms) plus processing overhead
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
 		defer cancel()
 
 		err = processor.Shutdown(shutdownCtx)

hitless/redis_connection_processor_test.go

@@ -368,6 +368,10 @@ func TestRedisConnectionProcessor(t *testing.T) {
 			if err := connections[i].MarkForHandoff("new-endpoint:6379", int64(i)); err != nil {
 				t.Fatalf("Failed to mark connection %d for handoff: %v", i, err)
 			}
+			// Set a mock initialization function
+			connections[i].SetInitConnFunc(func(ctx context.Context, cn *pool.Conn) error {
+				return nil
+			})
 		}
 
 		ctx := context.Background()
@@ -450,6 +454,11 @@ func TestRedisConnectionProcessor(t *testing.T) {
 			t.Fatalf("Failed to mark connection for handoff: %v", err)
 		}
 
+		// Set a mock initialization function
+		conn.SetInitConnFunc(func(ctx context.Context, cn *pool.Conn) error {
+			return nil
+		})
+
 		// Process the connection to trigger handoff
 		shouldPool, shouldRemove, err := processor.ProcessConnectionOnPut(ctx, conn)
 		if err != nil {
@@ -530,6 +539,11 @@ func TestRedisConnectionProcessor(t *testing.T) {
 			t.Fatalf("Failed to mark connection for handoff: %v", err)
 		}
 
+		// Set a mock initialization function
+		conn.SetInitConnFunc(func(ctx context.Context, cn *pool.Conn) error {
+			return nil
+		})
+
 		// Connection should no longer be usable
 		if conn.IsUsable() {
 			t.Error("Connection should not be usable after being marked for handoff")
@@ -598,6 +612,10 @@ func TestRedisConnectionProcessor(t *testing.T) {
 			if err := conn.MarkForHandoff("new-endpoint:6379", int64(i+1)); err != nil {
 				t.Fatalf("Failed to mark connection %d for handoff: %v", i, err)
 			}
+			// Set a mock initialization function
+			conn.SetInitConnFunc(func(ctx context.Context, cn *pool.Conn) error {
+				return nil
+			})
 
 			shouldPool, shouldRemove, err := processor.ProcessConnectionOnPut(ctx, conn)
 			if err != nil {
@@ -695,7 +713,14 @@ func TestRedisConnectionProcessor(t *testing.T) {
 		defer processor.Shutdown(context.Background())
 
 		conn := createMockPoolConnection()
-		conn.MarkForHandoff("new-endpoint:6379", 12345)
+		if err := conn.MarkForHandoff("new-endpoint:6379", 12345); err != nil {
+			t.Fatalf("Failed to mark connection for handoff: %v", err)
+		}
+
+		// Set a mock initialization function
+		conn.SetInitConnFunc(func(ctx context.Context, cn *pool.Conn) error {
+			return nil
+		})
 
 		ctx := context.Background()
 		shouldPool, shouldRemove, err := processor.ProcessConnectionOnPut(ctx, conn)

internal/pool/conn.go

@@ -323,6 +323,7 @@ func (cn *Conn) ExecuteInitConn(ctx context.Context) error {
 		}
 		cn.Inited = true
 		cn.setUsable(true) // Use atomic operation
+		return nil
 	}
 	return fmt.Errorf("redis: no initConnFunc set for connection %d", cn.GetID())
 }