From 8215f4126e7d788a99d3473ecfb217b44d94d093 Mon Sep 17 00:00:00 2001
From: Dom Dwyer <dom@itsallbroken.com>
Date: Tue, 24 Jan 2023 13:38:36 +0100
Subject: [PATCH] test: router balancer recovery

Ensure a recovering node is yielded from the balancer.
---
 router/src/dml_handlers/rpc_write/balancer.rs | 39 +++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/router/src/dml_handlers/rpc_write/balancer.rs b/router/src/dml_handlers/rpc_write/balancer.rs
index f421c8fafd..1ad1fc4d57 100644
--- a/router/src/dml_handlers/rpc_write/balancer.rs
+++ b/router/src/dml_handlers/rpc_write/balancer.rs
@@ -177,6 +177,45 @@ mod tests {
         assert_eq!(circuit_err_2.err_count(), 0);
     }
 
+    /// An unhealthy node that recovers is yielded to the caller.
+    #[tokio::test]
+    async fn test_balancer_upstream_recovery() {
+        const BALANCER_CALLS: usize = 10;
+
+        // Initialise 3 RPC clients and configure their mock circuit breakers;
+        // two returns a unhealthy state, one is healthy.
+        let circuit = Arc::new(MockCircuitBreaker::default());
+        circuit.set_usable(false);
+        let client = CircuitBreakingClient::new(Arc::new(MockWriteClient::default()))
+            .with_circuit_breaker(Arc::clone(&circuit));
+
+        assert_eq!(circuit.ok_count(), 0);
+
+        let balancer = Balancer::new([client]);
+
+        let mut endpoints = balancer.endpoints();
+        assert_matches!(endpoints.next(), None);
+        assert_eq!(circuit.is_usable_count(), 1);
+
+        circuit.set_usable(true);
+
+        let mut endpoints = balancer.endpoints();
+        assert_matches!(endpoints.next(), Some(_));
+        assert_eq!(circuit.is_usable_count(), 2);
+
+        // The now-healthy client is constantly yielded.
+        const N: usize = 3;
+        for _ in 0..N {
+            endpoints
+                .next()
+                .expect("should yield healthy client")
+                .write(WriteRequest::default())
+                .await
+                .expect("should succeed");
+        }
+        assert_eq!(circuit.ok_count(), N);
+    }
+
     // Ensure the balancer round-robins across all healthy clients.
     //
     // Note this is a property test that asserts the even distribution of the