ctp: move ChannelHandler into tests

Author: teskje

src/service/src/transport.rs

@@ -26,7 +26,7 @@ use std::convert::Infallible;
 use std::fmt::Debug;
 use std::time::Duration;
 
-use anyhow::{anyhow, bail};
+use anyhow::bail;
 use async_trait::async_trait;
 use bincode::Options;
 use futures::future;
@@ -407,38 +407,6 @@ impl<Out: Message, In: Message> Connection<Out, In> {
     }
 }
 
-/// A connection handler that simply forwards messages over channels.
-#[derive(Debug)]
-pub struct ChannelHandler<In, Out> {
-    tx: mpsc::UnboundedSender<In>,
-    rx: mpsc::UnboundedReceiver<Out>,
-}
-
-impl<In, Out> ChannelHandler<In, Out> {
-    pub fn new(tx: mpsc::UnboundedSender<In>, rx: mpsc::UnboundedReceiver<Out>) -> Self {
-        Self { tx, rx }
-    }
-}
-
-#[async_trait]
-impl<In: Message, Out: Message> GenericClient<In, Out> for ChannelHandler<In, Out> {
-    async fn send(&mut self, cmd: In) -> anyhow::Result<()> {
-        let result = self.tx.send(cmd);
-        result.map_err(|_| anyhow!("client channel disconnected"))
-    }
-
-    /// # Cancel safety
-    ///
-    /// This method is cancel safe.
-    async fn recv(&mut self) -> anyhow::Result<Option<Out>> {
-        // `mpsc::UnboundedReceiver::recv` is cancel safe.
-        match self.rx.recv().await {
-            Some(resp) => Ok(Some(resp)),
-            None => bail!("client channel disconnected"),
-        }
-    }
-}
-
 /// Perform the CTP handshake.
 ///
 /// To perform the handshake, each endpoint sends the protocol magic number, followed by a

src/service/tests/transport.rs

@@ -13,13 +13,14 @@ use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex, Once};
 use std::time::Duration;
 
+use anyhow::{anyhow, bail};
 use async_trait::async_trait;
 use futures::future;
 use mz_ore::assert_none;
 use mz_ore::netio::Listener;
 use mz_ore::retry::Retry;
 use mz_service::client::GenericClient;
-use mz_service::transport::{self, ChannelHandler, Message, NoopMetrics};
+use mz_service::transport::{self, Message, NoopMetrics};
 use semver::Version;
 use tokio::io::AsyncWriteExt;
 use tokio::sync::{mpsc, oneshot};
@@ -556,6 +557,38 @@ fn test_metrics() {
     sim.run().unwrap();
 }
 
+/// A connection handler that simply forwards messages over channels.
+#[derive(Debug)]
+pub struct ChannelHandler<In, Out> {
+    tx: mpsc::UnboundedSender<In>,
+    rx: mpsc::UnboundedReceiver<Out>,
+}
+
+impl<In, Out> ChannelHandler<In, Out> {
+    pub fn new(tx: mpsc::UnboundedSender<In>, rx: mpsc::UnboundedReceiver<Out>) -> Self {
+        Self { tx, rx }
+    }
+}
+
+#[async_trait]
+impl<In: Message, Out: Message> GenericClient<In, Out> for ChannelHandler<In, Out> {
+    async fn send(&mut self, cmd: In) -> anyhow::Result<()> {
+        let result = self.tx.send(cmd);
+        result.map_err(|_| anyhow!("client channel disconnected"))
+    }
+
+    /// # Cancel safety
+    ///
+    /// This method is cancel safe.
+    async fn recv(&mut self) -> anyhow::Result<Option<Out>> {
+        // `mpsc::Receiver::recv` is cancel safe.
+        match self.rx.recv().await {
+            Some(resp) => Ok(Some(resp)),
+            None => bail!("client channel disconnected"),
+        }
+    }
+}
+
 /// A connection handler that produces a single outbound message and then becomes silent.
 #[derive(Debug)]
 struct OneOutputHandler {