be able to reject connection OnGet

maintnotifications/pool_hook_test.go

@@ -360,10 +360,13 @@ func TestConnectionHook(t *testing.T) {
 		conn := createMockPoolConnection()
 
 		ctx := context.Background()
-		err := processor.OnGet(ctx, conn, false)
+		acceptCon, err := processor.OnGet(ctx, conn, false)
 		if err != nil {
 			t.Errorf("OnGet should not error for normal connection: %v", err)
 		}
+		if !acceptCon {
+			t.Error("Connection should be accepted for normal connection")
+		}
 	})
 
 	t.Run("OnGetWithPendingHandoff", func(t *testing.T) {
@@ -385,10 +388,13 @@ func TestConnectionHook(t *testing.T) {
 		conn.MarkQueuedForHandoff()                                 // Mark as queued (sets usable=false)
 
 		ctx := context.Background()
-		err := processor.OnGet(ctx, conn, false)
+		acceptCon, err := processor.OnGet(ctx, conn, false)
 		if err != ErrConnectionMarkedForHandoff {
 			t.Errorf("Expected ErrConnectionMarkedForHandoff, got %v", err)
 		}
+		if acceptCon {
+			t.Error("Connection should not be accepted when marked for handoff")
+		}
 
 		// Clean up
 		processor.GetPendingMap().Delete(conn)
@@ -416,10 +422,13 @@ func TestConnectionHook(t *testing.T) {
 
 		// Test OnGet with pending handoff
 		ctx := context.Background()
-		err := processor.OnGet(ctx, conn, false)
+		acceptCon, err := processor.OnGet(ctx, conn, false)
 		if err != ErrConnectionMarkedForHandoff {
 			t.Error("Should return ErrConnectionMarkedForHandoff for pending connection")
 		}
+		if acceptCon {
+			t.Error("Should not accept connection with pending handoff")
+		}
 
 		// Test removing from pending map and clearing handoff state
 		processor.GetPendingMap().Delete(conn)
@@ -432,10 +441,13 @@ func TestConnectionHook(t *testing.T) {
 		conn.SetUsable(true) // Make connection usable again
 
 		// Test OnGet without pending handoff
-		err = processor.OnGet(ctx, conn, false)
+		acceptCon, err = processor.OnGet(ctx, conn, false)
 		if err != nil {
 			t.Errorf("Should not return error for non-pending connection: %v", err)
 		}
+		if !acceptCon {
+			t.Error("Should accept connection without pending handoff")
+		}
 	})
 
 	t.Run("EventDrivenQueueOptimization", func(t *testing.T) {
@@ -628,11 +640,15 @@ func TestConnectionHook(t *testing.T) {
 		}
 
 		// OnGet should succeed for usable connection
-		err := processor.OnGet(ctx, conn, false)
+		acceptConn, err := processor.OnGet(ctx, conn, false)
 		if err != nil {
 			t.Errorf("OnGet should succeed for usable connection: %v", err)
 		}
 
+		if !acceptConn {
+			t.Error("Connection should be accepted when usable")
+		}
+
 		// Mark connection for handoff
 		if err := conn.MarkForHandoff("new-endpoint:6379", 1); err != nil {
 			t.Fatalf("Failed to mark connection for handoff: %v", err)
@@ -652,13 +668,17 @@ func TestConnectionHook(t *testing.T) {
 		}
 
 		// OnGet should fail for connection marked for handoff
-		err = processor.OnGet(ctx, conn, false)
+		acceptConn, err = processor.OnGet(ctx, conn, false)
 		if err == nil {
 			t.Error("OnGet should fail for connection marked for handoff")
 		}
+
 		if err != ErrConnectionMarkedForHandoff {
 			t.Errorf("Expected ErrConnectionMarkedForHandoff, got %v", err)
 		}
+		if acceptConn {
+			t.Error("Connection should not be accepted when marked for handoff")
+		}
 
 		// Process the connection to trigger handoff
 		shouldPool, shouldRemove, err := processor.OnPut(ctx, conn)
@@ -678,11 +698,15 @@ func TestConnectionHook(t *testing.T) {
 		}
 
 		// OnGet should succeed again
-		err = processor.OnGet(ctx, conn, false)
+		acceptConn, err = processor.OnGet(ctx, conn, false)
 		if err != nil {
 			t.Errorf("OnGet should succeed after handoff completion: %v", err)
 		}
 
+		if !acceptConn {
+			t.Error("Connection should be accepted after handoff completion")
+		}
+
 		t.Logf("Usable flag behavior test completed successfully")
 	})