From ff264e144defcfc1256a720caa73aa148ac750e1 Mon Sep 17 00:00:00 2001
From: KrishnaSindhur <krishnasindhur@gmail.com>
Date: Tue, 28 Jul 2026 08:39:11 +0530
Subject: [PATCH 1/2] fixed flakey test

---
 pgdog/src/backend/pool/lb/test.rs | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/pgdog/src/backend/pool/lb/test.rs b/pgdog/src/backend/pool/lb/test.rs
index f853e18df..d37b0a7a6 100644
--- a/pgdog/src/backend/pool/lb/test.rs
+++ b/pgdog/src/backend/pool/lb/test.rs
@@ -538,7 +538,7 @@ async fn test_read_write_split_exclude_primary_no_primary() {
     let replicas = LoadBalancer::new(
         &None,
         &replica_configs,
-        LoadBalancingStrategy::Random,
+        LoadBalancingStrategy::RoundRobin,
         ReadWriteSplit::ExcludePrimary,
     );
     replicas.launch();
@@ -546,13 +546,14 @@ async fn test_read_write_split_exclude_primary_no_primary() {
     let request = Request::default();
 
     // Should work normally with just replicas
+    let expected: HashSet<_> = replicas.pools().iter().map(|pool| pool.id()).collect();
     let mut replica_ids = HashSet::new();
-    for _ in 0..10 {
+    for _ in 0..expected.len() {
         let conn = replicas.get(&request).await.unwrap();
         replica_ids.insert(conn.pool.id());
     }
 
-    assert_eq!(replica_ids.len(), 2);
+    assert_eq!(replica_ids, expected);
 
     replicas.shutdown();
 }
@@ -568,7 +569,7 @@ async fn test_read_write_split_include_primary_no_primary() {
     let replicas = LoadBalancer::new(
         &None,
         &replica_configs,
-        LoadBalancingStrategy::Random,
+        LoadBalancingStrategy::RoundRobin,
         ReadWriteSplit::IncludePrimary,
     );
     replicas.launch();
@@ -576,14 +577,15 @@ async fn test_read_write_split_include_primary_no_primary() {
     let request = Request::default();
 
     // Should work normally with just replicas
+    let expected: HashSet<_> = replicas.pools().iter().map(|pool| pool.id()).collect();
     let mut replica_ids = HashSet::new();
-    for _ in 0..10 {
+    for _ in 0..expected.len() {
         let conn = replicas.get(&request).await.unwrap();
         replica_ids.insert(conn.pool.id());
     }
 
     // Should use both replica pools
-    assert_eq!(replica_ids.len(), 2);
+    assert_eq!(replica_ids, expected);
 
     replicas.shutdown();
 }

From 53ebe6b5433bc0dc556350d5d2735e272555a5cf Mon Sep 17 00:00:00 2001
From: KrishnaSindhur <krishnasindhur@gmail.com>
Date: Tue, 28 Jul 2026 22:30:02 +0530
Subject: [PATCH 2/2] =?UTF-8?q?fixed=20flakey=20test=20changing=20=2010=20?=
 =?UTF-8?q?=E2=86=92=20100?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pgdog/src/backend/pool/lb/test.rs | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/pgdog/src/backend/pool/lb/test.rs b/pgdog/src/backend/pool/lb/test.rs
index d37b0a7a6..f0d0dfb1b 100644
--- a/pgdog/src/backend/pool/lb/test.rs
+++ b/pgdog/src/backend/pool/lb/test.rs
@@ -538,7 +538,7 @@ async fn test_read_write_split_exclude_primary_no_primary() {
     let replicas = LoadBalancer::new(
         &None,
         &replica_configs,
-        LoadBalancingStrategy::RoundRobin,
+        LoadBalancingStrategy::Random,
         ReadWriteSplit::ExcludePrimary,
     );
     replicas.launch();
@@ -546,14 +546,13 @@ async fn test_read_write_split_exclude_primary_no_primary() {
     let request = Request::default();
 
     // Should work normally with just replicas
-    let expected: HashSet<_> = replicas.pools().iter().map(|pool| pool.id()).collect();
     let mut replica_ids = HashSet::new();
-    for _ in 0..expected.len() {
+    for _ in 0..100 {
         let conn = replicas.get(&request).await.unwrap();
         replica_ids.insert(conn.pool.id());
     }
 
-    assert_eq!(replica_ids, expected);
+    assert_eq!(replica_ids.len(), 2);
 
     replicas.shutdown();
 }
@@ -569,7 +568,7 @@ async fn test_read_write_split_include_primary_no_primary() {
     let replicas = LoadBalancer::new(
         &None,
         &replica_configs,
-        LoadBalancingStrategy::RoundRobin,
+        LoadBalancingStrategy::Random,
         ReadWriteSplit::IncludePrimary,
     );
     replicas.launch();
@@ -577,15 +576,14 @@ async fn test_read_write_split_include_primary_no_primary() {
     let request = Request::default();
 
     // Should work normally with just replicas
-    let expected: HashSet<_> = replicas.pools().iter().map(|pool| pool.id()).collect();
     let mut replica_ids = HashSet::new();
-    for _ in 0..expected.len() {
+    for _ in 0..100 {
         let conn = replicas.get(&request).await.unwrap();
         replica_ids.insert(conn.pool.id());
     }
 
     // Should use both replica pools
-    assert_eq!(replica_ids, expected);
+    assert_eq!(replica_ids.len(), 2);
 
     replicas.shutdown();
 }