grpc: add testing utilities for LB policy tests

grpc/src/client/load_balancing/mod.rs

@@ -53,6 +53,8 @@ use crate::client::{
 
 pub mod child_manager;
 pub mod pick_first;
+#[cfg(test)]
+pub mod test_utils;
 
 pub(crate) mod registry;
 use super::{service_config::LbConfig, subchannel::SubchannelStateWatcher};

grpc/src/client/load_balancing/test_utils.rs

@@ -0,0 +1,147 @@
+/*
+ *
+ * Copyright 2025 gRPC authors.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+ * IN THE SOFTWARE.
+ *
+ */
+
+use crate::client::load_balancing::{
+    ChannelController, ExternalSubchannel, ForwardingSubchannel, LbState, Subchannel, WorkScheduler,
+};
+use crate::client::name_resolution::Address;
+use crate::service::{Message, Request, Response, Service};
+use std::hash::{Hash, Hasher};
+use std::{fmt::Debug, ops::Add, sync::Arc};
+use tokio::sync::{mpsc, Notify};
+use tokio::task::AbortHandle;
+
+pub(crate) struct EmptyMessage {}
+impl Message for EmptyMessage {}
+pub(crate) fn new_request() -> Request {
+    Request::new(Box::pin(tokio_stream::once(
+        Box::new(EmptyMessage {}) as Box<dyn Message>
+    )))
+}
+
+// A test subchannel that forwards connect calls to a channel.
+// This allows tests to verify when a subchannel is asked to connect.
+pub(crate) struct TestSubchannel {
+    address: Address,
+    tx_connect: mpsc::UnboundedSender<TestEvent>,
+}
+
+impl TestSubchannel {
+    fn new(address: Address, tx_connect: mpsc::UnboundedSender<TestEvent>) -> Self {
+        Self {
+            address,
+            tx_connect,
+        }
+    }
+}
+
+impl ForwardingSubchannel for TestSubchannel {
+    fn delegate(&self) -> Arc<dyn Subchannel> {
+        panic!("unsupported operation on a test subchannel");
+    }
+
+    fn address(&self) -> Address {
+        self.address.clone()
+    }
+
+    fn connect(&self) {
+        println!("connect called for subchannel {}", self.address);
+        self.tx_connect
+            .send(TestEvent::Connect(self.address.clone()))
+            .unwrap();
+    }
+}
+
+impl Hash for TestSubchannel {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.address.hash(state);
+    }
+}
+
+impl PartialEq for TestSubchannel {
+    fn eq(&self, other: &Self) -> bool {
+        std::ptr::eq(self, other)
+    }
+}
+impl Eq for TestSubchannel {}
+
+pub(crate) enum TestEvent {
+    NewSubchannel(Arc<dyn Subchannel>),
+    UpdatePicker(LbState),
+    RequestResolution,
+    Connect(Address),
+    ScheduleWork,
+}
+
+// TODO(easwars): Remove this and instead derive Debug.
+impl Debug for TestEvent {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::NewSubchannel(sc) => write!(f, "NewSubchannel({})", sc.address()),
+            Self::UpdatePicker(state) => write!(f, "UpdatePicker({})", state.connectivity_state),
+            Self::RequestResolution => write!(f, "RequestResolution"),
+            Self::Connect(addr) => write!(f, "Connect({})", addr.address.to_string()),
+            Self::ScheduleWork => write!(f, "ScheduleWork"),
+        }
+    }
+}
+
+/// A test channel controller that forwards calls to a channel.  This allows
+/// tests to verify when a channel controller is asked to create subchannels or
+/// update the picker.
+pub(crate) struct TestChannelController {
+    pub(crate) tx_events: mpsc::UnboundedSender<TestEvent>,
+}
+
+impl ChannelController for TestChannelController {
+    fn new_subchannel(&mut self, address: &Address) -> Arc<dyn Subchannel> {
+        println!("new_subchannel called for address {}", address);
+        let notify = Arc::new(Notify::new());
+        let subchannel: Arc<dyn Subchannel> =
+            Arc::new(TestSubchannel::new(address.clone(), self.tx_events.clone()));
+        self.tx_events
+            .send(TestEvent::NewSubchannel(subchannel.clone()))
+            .unwrap();
+        subchannel
+    }
+    fn update_picker(&mut self, update: LbState) {
+        println!("picker_update called with {}", update.connectivity_state);
+        self.tx_events
+            .send(TestEvent::UpdatePicker(update))
+            .unwrap();
+    }
+    fn request_resolution(&mut self) {
+        self.tx_events.send(TestEvent::RequestResolution).unwrap();
+    }
+}
+
+pub(crate) struct TestWorkScheduler {
+    pub(crate) tx_events: mpsc::UnboundedSender<TestEvent>,
+}
+
+impl WorkScheduler for TestWorkScheduler {
+    fn schedule_work(&self) {
+        self.tx_events.send(TestEvent::ScheduleWork).unwrap();
+    }
+}