Replace gRPC Backend (#10)

**Rationale:**

Having two separate servers and communication methods resulted in additional maintenance & the need to convert often between backend & frontend data types.
By moving the backend communication off of gRPC and to just use websockets it both gives more control & allows for simplification of the implementation.

#8

**Changes:**

- Replaces gRPC backend.
  - New implementation automatically handles reconnect logic
- Implements an api layer
- Migrates examples to the api layer
- Implements a proc macro to make command handling easier
- Implements unit tests for the api layer (90+% coverage)
- Implements integration tests for the proc macro (90+% coverage)

Reviewed-on: #10
Co-authored-by: Sergey Savelyev <sergeysav.nn@gmail.com>
Co-committed-by: Sergey Savelyev <sergeysav.nn@gmail.com>

api/Cargo.toml

@@ -0,0 +1,24 @@
+
+[package]
+name = "api"
+edition = "2021"
+version = "0.1.0"
+authors = ["Sergey <me@sergeysav.com>"]
+
+[dependencies]
+api-core = { path = "../api-core" }
+api-proc-macro = { path = "../api-proc-macro" }
+chrono = { workspace = true, features = ["serde"] }
+derive_more = { workspace = true, features = ["from", "try_into"] }
+futures-util = { workspace = true }
+log = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+serde_json = { workspace = true }
+thiserror = { workspace = true }
+tokio = { workspace = true, features = ["rt", "macros", "time"] }
+tokio-tungstenite = { workspace = true, features = ["rustls-tls-native-roots"] }
+tokio-util = { workspace = true }
+uuid = { workspace = true, features = ["serde"] }
+
+[dev-dependencies]
+env_logger = { workspace = true }

api/src/client/command.rs

@@ -0,0 +1,454 @@
+use crate::client::Client;
+use crate::messages::command::CommandResponse;
+use api_core::command::{CommandHeader, IntoCommandDefinition};
+use std::fmt::Display;
+use std::sync::Arc;
+use tokio::select;
+use tokio_util::sync::CancellationToken;
+
+pub struct CommandRegistry {
+    client: Arc<Client>,
+}
+
+impl CommandRegistry {
+    pub fn new(client: Arc<Client>) -> Self {
+        Self { client }
+    }
+
+    pub fn register_handler<C: IntoCommandDefinition, F, E: Display>(
+        &self,
+        command_name: impl Into<String>,
+        mut callback: F,
+    ) -> CommandHandle
+    where
+        F: FnMut(CommandHeader, C) -> Result<String, E> + Send + 'static,
+    {
+        let cancellation_token = CancellationToken::new();
+        let result = CommandHandle {
+            cancellation_token: cancellation_token.clone(),
+        };
+        let client = self.client.clone();
+
+        let command_definition = C::create(command_name.into());
+
+        tokio::spawn(async move {
+            while !cancellation_token.is_cancelled() {
+                // This would only fail if the sender closed while trying to insert data
+                // It would wait until space is made
+                let Ok(mut rx) = client
+                    .register_callback_channel(command_definition.clone())
+                    .await
+                else {
+                    continue;
+                };
+
+                loop {
+                    // select used so that this loop gets broken if the token is cancelled
+                    select!(
+                        rx_value = rx.recv() => {
+                            if let Some((cmd, responder)) = rx_value {
+                                let header = cmd.header.clone();
+                                let response = match C::parse(cmd) {
+                                    Ok(cmd) => match callback(header, cmd) {
+                                        Ok(response) => CommandResponse {
+                                            success: true,
+                                            response,
+                                        },
+                                        Err(err) => CommandResponse {
+                                            success: false,
+                                            response: err.to_string(),
+                                        },
+                                    },
+                                    Err(err) => CommandResponse {
+                                        success: false,
+                                        response: err.to_string(),
+                                    },
+                                };
+                                // This should only err if we had an error elsewhere
+                                let _ = responder.send(response);
+                            } else {
+                                break;
+                            }
+                        },
+                        _ = cancellation_token.cancelled() => { break; },
+                    );
+                }
+            }
+        });
+
+        result
+    }
+}
+
+pub struct CommandHandle {
+    cancellation_token: CancellationToken,
+}
+
+impl Drop for CommandHandle {
+    fn drop(&mut self) {
+        self.cancellation_token.cancel();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::client::command::CommandRegistry;
+    use crate::client::tests::create_test_client;
+    use crate::client::Callback;
+    use crate::messages::callback::GenericCallbackError;
+    use crate::messages::command::CommandResponse;
+    use crate::messages::payload::RequestMessagePayload;
+    use crate::messages::telemetry_definition::TelemetryDefinitionResponse;
+    use crate::messages::ResponseMessage;
+    use api_core::command::{
+        Command, CommandDefinition, CommandHeader, CommandParameterDefinition,
+        IntoCommandDefinition, IntoCommandDefinitionError,
+    };
+    use api_core::data_type::DataType;
+    use std::collections::HashMap;
+    use std::convert::Infallible;
+    use std::sync::Arc;
+    use std::time::Duration;
+    use tokio::sync::oneshot;
+    use tokio::time::timeout;
+    use uuid::Uuid;
+
+    struct CmdType {
+        #[allow(unused)]
+        param1: f32,
+    }
+
+    impl IntoCommandDefinition for CmdType {
+        fn create(name: String) -> CommandDefinition {
+            CommandDefinition {
+                name,
+                parameters: vec![CommandParameterDefinition {
+                    name: "param1".to_string(),
+                    data_type: DataType::Float32,
+                }],
+            }
+        }
+
+        fn parse(command: Command) -> Result<Self, IntoCommandDefinitionError> {
+            Ok(Self {
+                param1: (*command.parameters.get("param1").ok_or_else(|| {
+                    IntoCommandDefinitionError::ParameterMissing("param1".to_string())
+                })?)
+                .try_into()
+                .map_err(|_| IntoCommandDefinitionError::MismatchedType {
+                    parameter: "param1".to_string(),
+                    expected: DataType::Float32,
+                })?,
+            })
+        }
+    }
+
+    #[tokio::test]
+    async fn simple_handler() {
+        // if _c drops then we are disconnected
+        let (mut rx, _c, client) = create_test_client();
+
+        let cmd_reg = CommandRegistry::new(Arc::new(client));
+
+        let _cmd_handle = cmd_reg.register_handler("cmd", |_, _: CmdType| {
+            Ok("success".to_string()) as Result<_, Infallible>
+        });
+
+        let msg = timeout(Duration::from_secs(1), rx.recv())
+            .await
+            .unwrap()
+            .unwrap();
+        let Callback::Registered(callback) = msg.callback else {
+            panic!("Incorrect Callback Type");
+        };
+
+        let mut params = HashMap::new();
+        params.insert("param1".to_string(), 0.0f32.into());
+
+        let (response_tx, response_rx) = oneshot::channel();
+        timeout(
+            Duration::from_secs(1),
+            callback.send((
+                ResponseMessage {
+                    uuid: Uuid::new_v4(),
+                    response: Some(msg.msg.uuid),
+                    payload: Command {
+                        header: CommandHeader {
+                            timestamp: Default::default(),
+                        },
+                        parameters: params,
+                    }
+                    .into(),
+                },
+                response_tx,
+            )),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+        let response = timeout(Duration::from_secs(1), response_rx)
+            .await
+            .unwrap()
+            .unwrap();
+        let RequestMessagePayload::CommandResponse(CommandResponse { success, response }) =
+            response
+        else {
+            panic!("Unexpected Response Type");
+        };
+        assert!(success);
+        assert_eq!(response, "success");
+    }
+
+    #[tokio::test]
+    async fn handler_failed() {
+        // if _c drops then we are disconnected
+        let (mut rx, _c, client) = create_test_client();
+
+        let cmd_reg = CommandRegistry::new(Arc::new(client));
+
+        let _cmd_handle = cmd_reg.register_handler("cmd", |_, _: CmdType| {
+            Err("failure".into()) as Result<_, Box<dyn std::error::Error>>
+        });
+
+        let msg = timeout(Duration::from_secs(1), rx.recv())
+            .await
+            .unwrap()
+            .unwrap();
+        let Callback::Registered(callback) = msg.callback else {
+            panic!("Incorrect Callback Type");
+        };
+
+        let mut params = HashMap::new();
+        params.insert("param1".to_string(), 1.0f32.into());
+
+        let (response_tx, response_rx) = oneshot::channel();
+        timeout(
+            Duration::from_secs(1),
+            callback.send((
+                ResponseMessage {
+                    uuid: Uuid::new_v4(),
+                    response: Some(msg.msg.uuid),
+                    payload: Command {
+                        header: CommandHeader {
+                            timestamp: Default::default(),
+                        },
+                        parameters: params,
+                    }
+                    .into(),
+                },
+                response_tx,
+            )),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+        let response = timeout(Duration::from_secs(1), response_rx)
+            .await
+            .unwrap()
+            .unwrap();
+        let RequestMessagePayload::CommandResponse(CommandResponse { success, response }) =
+            response
+        else {
+            panic!("Unexpected Response Type");
+        };
+        assert!(!success);
+        assert_eq!(response, "failure");
+    }
+
+    #[tokio::test]
+    async fn parse_failed() {
+        // if _c drops then we are disconnected
+        let (mut rx, _c, client) = create_test_client();
+
+        let cmd_reg = CommandRegistry::new(Arc::new(client));
+
+        let _cmd_handle = cmd_reg.register_handler("cmd", |_, _: CmdType| {
+            Err("failure".into()) as Result<_, Box<dyn std::error::Error>>
+        });
+
+        let msg = timeout(Duration::from_secs(1), rx.recv())
+            .await
+            .unwrap()
+            .unwrap();
+        let Callback::Registered(callback) = msg.callback else {
+            panic!("Incorrect Callback Type");
+        };
+
+        let mut params = HashMap::new();
+        params.insert("param1".to_string(), 1.0f64.into());
+
+        let (response_tx, response_rx) = oneshot::channel();
+        timeout(
+            Duration::from_secs(1),
+            callback.send((
+                ResponseMessage {
+                    uuid: Uuid::new_v4(),
+                    response: Some(msg.msg.uuid),
+                    payload: Command {
+                        header: CommandHeader {
+                            timestamp: Default::default(),
+                        },
+                        parameters: params,
+                    }
+                    .into(),
+                },
+                response_tx,
+            )),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+        let response = timeout(Duration::from_secs(1), response_rx)
+            .await
+            .unwrap()
+            .unwrap();
+        let RequestMessagePayload::CommandResponse(CommandResponse {
+            success,
+            response: _,
+        }) = response
+        else {
+            panic!("Unexpected Response Type");
+        };
+        assert!(!success);
+    }
+
+    #[tokio::test]
+    async fn wrong_message() {
+        // if _c drops then we are disconnected
+        let (mut rx, _c, client) = create_test_client();
+
+        let cmd_reg = CommandRegistry::new(Arc::new(client));
+
+        let _cmd_handle =
+            cmd_reg.register_handler("cmd", |_, _: CmdType| -> Result<_, Infallible> {
+                panic!("This should not happen");
+            });
+
+        let msg = timeout(Duration::from_secs(1), rx.recv())
+            .await
+            .unwrap()
+            .unwrap();
+        let Callback::Registered(callback) = msg.callback else {
+            panic!("Incorrect Callback Type");
+        };
+
+        let (response_tx, response_rx) = oneshot::channel();
+        timeout(
+            Duration::from_secs(1),
+            callback.send((
+                ResponseMessage {
+                    uuid: Uuid::new_v4(),
+                    response: Some(msg.msg.uuid),
+                    payload: TelemetryDefinitionResponse {
+                        uuid: Uuid::new_v4(),
+                    }
+                    .into(),
+                },
+                response_tx,
+            )),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+        let response = timeout(Duration::from_secs(1), response_rx)
+            .await
+            .unwrap()
+            .unwrap();
+        let RequestMessagePayload::GenericCallbackError(err) = response else {
+            panic!("Unexpected Response Type");
+        };
+        assert_eq!(err, GenericCallbackError::MismatchedType);
+    }
+
+    #[tokio::test]
+    async fn callback_closed() {
+        // if _c drops then we are disconnected
+        let (mut rx, _c, client) = create_test_client();
+
+        let cmd_reg = CommandRegistry::new(Arc::new(client));
+
+        let cmd_handle =
+            cmd_reg.register_handler("cmd", |_, _: CmdType| -> Result<_, Infallible> {
+                panic!("This should not happen");
+            });
+
+        let msg = timeout(Duration::from_secs(1), rx.recv())
+            .await
+            .unwrap()
+            .unwrap();
+        let Callback::Registered(callback) = msg.callback else {
+            panic!("Incorrect Callback Type");
+        };
+
+        // This should shut down the command handler
+        drop(cmd_handle);
+
+        // Send a command
+        let mut params = HashMap::new();
+        params.insert("param1".to_string(), 0.0f32.into());
+        let (response_tx, response_rx) = oneshot::channel();
+        timeout(
+            Duration::from_secs(1),
+            callback.send((
+                ResponseMessage {
+                    uuid: Uuid::new_v4(),
+                    response: Some(msg.msg.uuid),
+                    payload: Command {
+                        header: CommandHeader {
+                            timestamp: Default::default(),
+                        },
+                        parameters: params,
+                    }
+                    .into(),
+                },
+                response_tx,
+            )),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+
+        let response = timeout(Duration::from_secs(1), response_rx)
+            .await
+            .unwrap()
+            .unwrap();
+        let RequestMessagePayload::GenericCallbackError(err) = response else {
+            panic!("Unexpected Response Type");
+        };
+        assert_eq!(err, GenericCallbackError::CallbackClosed);
+    }
+
+    #[tokio::test]
+    async fn reconnect() {
+        // if _c drops then we are disconnected
+        let (mut rx, _c, client) = create_test_client();
+
+        let cmd_reg = CommandRegistry::new(Arc::new(client));
+
+        let _cmd_handle =
+            cmd_reg.register_handler("cmd", |_, _: CmdType| -> Result<_, Infallible> {
+                panic!("This should not happen");
+            });
+
+        let msg = timeout(Duration::from_secs(1), rx.recv())
+            .await
+            .unwrap()
+            .unwrap();
+        let Callback::Registered(callback) = msg.callback else {
+            panic!("Incorrect Callback Type");
+        };
+
+        println!("Dropping");
+        drop(callback);
+        println!("Dropped");
+
+        // The command re-registers itself
+        let msg = timeout(Duration::from_secs(1), rx.recv())
+            .await
+            .unwrap()
+            .unwrap();
+        let Callback::Registered(_) = msg.callback else {
+            panic!("Incorrect Callback Type");
+        };
+    }
+}

api/src/client/config.rs

@@ -0,0 +1,11 @@
+pub struct ClientConfiguration {
+    pub send_buffer_size: usize,
+}
+
+impl Default for ClientConfiguration {
+    fn default() -> Self {
+        Self {
+            send_buffer_size: 128,
+        }
+    }
+}

api/src/client/context.rs

@@ -0,0 +1,594 @@
+use crate::client::config::ClientConfiguration;
+use crate::client::error::{ConnectError, MessageError};
+use crate::client::{Callback, ClientChannel, OutgoingMessage, RegisteredCallback};
+use crate::messages::callback::GenericCallbackError;
+use crate::messages::payload::RequestMessagePayload;
+use crate::messages::{RequestMessage, ResponseMessage};
+use futures_util::{Sink, SinkExt, Stream, StreamExt};
+use log::{debug, error, info, trace, warn};
+use std::collections::HashMap;
+use std::fmt::Display;
+use std::sync::mpsc::sync_channel;
+use std::thread;
+use std::time::Duration;
+use tokio::sync::{mpsc, oneshot, watch, RwLockWriteGuard};
+use tokio::time::sleep;
+use tokio::{select, spawn};
+use tokio_tungstenite::connect_async;
+use tokio_tungstenite::tungstenite::handshake::client::{Request, Response as TungResponse};
+use tokio_tungstenite::tungstenite::{Error as TungError, Message};
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+pub struct ClientContext {
+    pub cancel: CancellationToken,
+    pub request: Request,
+    pub connected_state_tx: watch::Sender<bool>,
+    pub client_configuration: ClientConfiguration,
+}
+
+impl ClientContext {
+    pub fn start(mut self, channel: ClientChannel) -> Result<(), ConnectError> {
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .build()?;
+
+        let (tx, rx) = sync_channel::<()>(1);
+
+        let _detached = thread::Builder::new()
+            .name("tlm-client".to_string())
+            .spawn(move || {
+                runtime.block_on(async {
+                    let mut write_lock = channel.write().await;
+
+                    // This cannot fail
+                    let _ = tx.send(());
+
+                    while !self.cancel.is_cancelled() {
+                        write_lock = self
+                            .run_connection(write_lock, &channel, connect_async)
+                            .await;
+                    }
+                    drop(write_lock);
+                });
+            })?;
+
+        // This cannot fail
+        let _ = rx.recv();
+
+        Ok(())
+    }
+
+    async fn run_connection<'a, F, W, E>(
+        &mut self,
+        mut write_lock: RwLockWriteGuard<'a, mpsc::Sender<OutgoingMessage>>,
+        channel: &'a ClientChannel,
+        mut connection_fn: F,
+    ) -> RwLockWriteGuard<'a, mpsc::Sender<OutgoingMessage>>
+    where
+        F: AsyncFnMut(Request) -> Result<(W, TungResponse), TungError>,
+        W: Stream<Item = Result<Message, TungError>> + Sink<Message, Error = E> + Unpin,
+        E: Display,
+    {
+        debug!("Attempting to Connect to {}", self.request.uri());
+        let mut ws = match connection_fn(self.request.clone()).await {
+            Ok((ws, _)) => ws,
+            Err(e) => {
+                info!("Failed to Connect: {e}");
+                sleep(Duration::from_secs(1)).await;
+                return write_lock;
+            }
+        };
+        info!("Connected to {}", self.request.uri());
+
+        let (tx, rx) = mpsc::channel(self.client_configuration.send_buffer_size);
+        *write_lock = tx;
+        drop(write_lock);
+
+        // Don't care about the previous value
+        let _ = self.connected_state_tx.send_replace(true);
+
+        let close_connection = self.handle_connection(&mut ws, rx, channel).await;
+
+        let write_lock = channel.write().await;
+        // Send this after grabbing the lock - to prevent extra contention when others try to grab
+        // the lock to use that as a signal that we have reconnected
+        let _ = self.connected_state_tx.send_replace(false);
+        if close_connection {
+            // Manually close to allow the impl trait to be used
+            if let Err(e) = ws.send(Message::Close(None)).await {
+                error!("Failed to Close the Connection: {e}");
+            }
+        }
+        write_lock
+    }
+
+    async fn handle_connection<W>(
+        &mut self,
+        ws: &mut W,
+        mut rx: mpsc::Receiver<OutgoingMessage>,
+        channel: &ClientChannel,
+    ) -> bool
+    where
+        W: Stream<Item = Result<Message, TungError>> + Sink<Message> + Unpin,
+        <W as Sink<Message>>::Error: Display,
+    {
+        let mut callbacks = HashMap::<Uuid, Callback>::new();
+        loop {
+            select! {
+                _ = self.cancel.cancelled() => { break; },
+                Some(msg) = ws.next() => {
+                    match msg {
+                        Ok(msg) => {
+                            match msg {
+                                Message::Text(msg) => {
+                                    trace!("Incoming: {msg}");
+                                    let msg: ResponseMessage = match serde_json::from_str(&msg) {
+                                        Ok(m) => m,
+                                        Err(e) => {
+                                            error!("Failed to deserialize {e}");
+                                            break;
+                                        }
+                                    };
+                                    self.handle_incoming(msg, &mut callbacks, channel).await;
+                                }
+                                Message::Binary(_) => unimplemented!("Binary Data Not Implemented"),
+                                Message::Ping(data) => {
+                                    if let Err(e) = ws.send(Message::Pong(data)).await {
+                                        error!("Failed to send Pong {e}");
+                                        break;
+                                    }
+                                }
+                                Message::Pong(_) => {
+                                    // Intentionally Left Empty
+                                }
+                                Message::Close(_) => {
+                                    debug!("Websocket Closed");
+                                    return false;
+                                }
+                                Message::Frame(_) => unreachable!("Not Possible"),
+                            }
+                        }
+                        Err(e) => {
+                            error!("Receive Error {e}");
+                            break;
+                        }
+                    }
+                }
+                Some(msg) = rx.recv() => {
+                    // Insert a callback if it isn't a None callback
+                    if !matches!(msg.callback, Callback::None) {
+                        callbacks.insert(msg.msg.uuid, msg.callback);
+                    }
+                    let msg = match serde_json::to_string(&msg.msg) {
+                        Ok(m) => m,
+                        Err(e) => {
+                            error!("Encode Error {e}");
+                            break;
+                        }
+                    };
+                    trace!("Outgoing: {msg}");
+                    if let Err(e) = ws.send(Message::Text(msg.into())).await {
+                        error!("Send Error {e}");
+                        break;
+                    }
+                }
+                else => { break; },
+            }
+        }
+        true
+    }
+
+    async fn handle_incoming(
+        &mut self,
+        msg: ResponseMessage,
+        callbacks: &mut HashMap<Uuid, Callback>,
+        channel: &ClientChannel,
+    ) {
+        if let Some(response_uuid) = msg.response {
+            match callbacks.get(&response_uuid) {
+                Some(Callback::None) => {
+                    callbacks.remove(&response_uuid);
+                    unreachable!("We skip registering callbacks of None type");
+                }
+                Some(Callback::Once(_)) => {
+                    let Some(Callback::Once(callback)) = callbacks.remove(&response_uuid) else {
+                        return;
+                    };
+                    let _ = callback.send(msg);
+                }
+                Some(Callback::Registered(callback)) => {
+                    let callback = callback.clone();
+                    spawn(Self::handle_registered_callback(
+                        callback,
+                        msg,
+                        channel.clone(),
+                    ));
+                }
+                None => {
+                    warn!("No Callback Registered for {response_uuid}");
+                }
+            }
+        }
+    }
+
+    async fn handle_registered_callback(
+        callback: RegisteredCallback,
+        msg: ResponseMessage,
+        channel: ClientChannel,
+    ) {
+        let (tx, rx) = oneshot::channel();
+
+        let uuid = msg.uuid;
+
+        let response = match callback.send((msg, tx)).await {
+            Err(_) => GenericCallbackError::CallbackClosed.into(),
+            Ok(()) => rx
+                .await
+                .unwrap_or_else(|_| GenericCallbackError::CallbackClosed.into()),
+        };
+
+        if let Err(e) = Self::send_response(channel, response, uuid).await {
+            error!("Failed to send response {e}");
+        }
+    }
+
+    async fn send_response(
+        channel: ClientChannel,
+        payload: RequestMessagePayload,
+        response_uuid: Uuid,
+    ) -> Result<(), MessageError> {
+        // If this failed that means we're in the middle of reconnecting, so our callbacks
+        // are all being cleaned up as-is. No response needed.
+        let sender = channel.try_read()?;
+        let data = sender.reserve().await?;
+
+        data.send(OutgoingMessage {
+            msg: RequestMessage {
+                uuid: Uuid::new_v4(),
+                response: Some(response_uuid),
+                payload,
+            },
+            callback: Callback::None,
+        });
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::messages::telemetry_definition::{
+        TelemetryDefinitionRequest, TelemetryDefinitionResponse,
+    };
+    use crate::test::mock_stream_sink::{create_mock_stream_sink, MockStreamSinkControl};
+    use api_core::data_type::DataType;
+    use log::LevelFilter;
+    use std::future::Future;
+    use std::ops::Deref;
+    use tokio::sync::mpsc::Sender;
+    use tokio::sync::RwLock;
+    use tokio::time::timeout;
+    use tokio::try_join;
+    use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+    use tokio_util::bytes::Bytes;
+
+    async fn assert_client_interaction<F, R>(future: F)
+    where
+        F: Send
+            + FnOnce(
+                Sender<OutgoingMessage>,
+                MockStreamSinkControl<Result<Message, TungError>, Message>,
+                CancellationToken,
+            ) -> R
+            + 'static,
+        R: Future<Output = ()> + Send,
+    {
+        let (control, stream_sink) =
+            create_mock_stream_sink::<Result<Message, TungError>, Message>();
+
+        let cancel_token = CancellationToken::new();
+        let inner_cancel_token = cancel_token.clone();
+        let (connected_state_tx, _connected_state_rx) = watch::channel(false);
+
+        let mut context = ClientContext {
+            cancel: cancel_token,
+            request: "mock".into_client_request().unwrap(),
+            connected_state_tx,
+            client_configuration: Default::default(),
+        };
+
+        let (tx, _rx) = mpsc::channel(1);
+        let channel = ClientChannel::new(RwLock::new(tx));
+        let used_channel = channel.clone();
+
+        let write_lock = used_channel.write().await;
+
+        let handle = spawn(async move {
+            let channel = channel;
+            let read = channel.read().await;
+            let sender = read.deref().clone();
+            drop(read);
+            future(sender, control, inner_cancel_token).await;
+        });
+
+        let mut stream_sink = Some(stream_sink);
+
+        let connection_fn = async |_: Request| {
+            let stream_sink = stream_sink.take().ok_or(TungError::ConnectionClosed)?;
+
+            Ok((stream_sink, TungResponse::default())) as Result<(_, _), TungError>
+        };
+
+        let context_result = async {
+            drop(
+                context
+                    .run_connection(write_lock, &used_channel, connection_fn)
+                    .await,
+            );
+            Ok(())
+        };
+
+        try_join!(context_result, timeout(Duration::from_secs(1), handle),)
+            .unwrap()
+            .1
+            .unwrap();
+    }
+
+    #[tokio::test]
+    async fn connection_closes_when_websocket_closes() {
+        let _ = env_logger::builder()
+            .is_test(true)
+            .filter_level(LevelFilter::Trace)
+            .try_init();
+
+        assert_client_interaction(|sender, mut control, _| async move {
+            let msg = Uuid::new_v4();
+            sender
+                .send(OutgoingMessage {
+                    msg: RequestMessage {
+                        uuid: msg,
+                        response: None,
+                        payload: TelemetryDefinitionRequest {
+                            name: "".to_string(),
+                            data_type: DataType::Float32,
+                        }
+                        .into(),
+                    },
+                    callback: Callback::None,
+                })
+                .await
+                .unwrap();
+            // We expect an outgoing message
+            assert!(matches!(
+                control.outgoing.recv().await.unwrap(),
+                Message::Text(_)
+            ));
+            // We receive an incoming close message
+            control
+                .incoming
+                .send(Ok(Message::Close(None)))
+                .await
+                .unwrap();
+            // Then we expect the outgoing to close with no message
+            assert!(control.outgoing.recv().await.is_none());
+            assert!(control.incoming.is_closed());
+        })
+        .await;
+    }
+
+    #[tokio::test]
+    async fn connection_closes_when_cancelled() {
+        let _ = env_logger::builder()
+            .is_test(true)
+            .filter_level(LevelFilter::Trace)
+            .try_init();
+
+        assert_client_interaction(|_, mut control, cancel| async move {
+            cancel.cancel();
+            // We expect an outgoing cancel message
+            assert!(matches!(
+                control.outgoing.recv().await.unwrap(),
+                Message::Close(_)
+            ));
+            // Then we expect to close with no message
+            assert!(control.outgoing.recv().await.is_none());
+            assert!(control.incoming.is_closed());
+        })
+        .await;
+    }
+
+    #[tokio::test]
+    async fn callback_request() {
+        let _ = env_logger::builder()
+            .is_test(true)
+            .filter_level(LevelFilter::Trace)
+            .try_init();
+
+        assert_client_interaction(|sender, mut control, _| async move {
+            let (callback_tx, callback_rx) = oneshot::channel();
+            let msg = Uuid::new_v4();
+            sender
+                .send(OutgoingMessage {
+                    msg: RequestMessage {
+                        uuid: msg,
+                        response: None,
+                        payload: TelemetryDefinitionRequest {
+                            name: "".to_string(),
+                            data_type: DataType::Float32,
+                        }
+                        .into(),
+                    },
+                    callback: Callback::Once(callback_tx),
+                })
+                .await
+                .unwrap();
+
+            // We expect an outgoing message
+            assert!(matches!(
+                control.outgoing.recv().await.unwrap(),
+                Message::Text(_)
+            ));
+
+            // Then we get an incoming message for this callback
+            let response_message = ResponseMessage {
+                uuid: Uuid::new_v4(),
+                response: Some(msg),
+                payload: TelemetryDefinitionResponse {
+                    uuid: Uuid::new_v4(),
+                }
+                .into(),
+            };
+            control
+                .incoming
+                .send(Ok(Message::Text(
+                    serde_json::to_string(&response_message).unwrap().into(),
+                )))
+                .await
+                .unwrap();
+
+            // We expect the callback to run
+            let message = callback_rx.await.unwrap();
+            // And give us the message we provided it
+            assert_eq!(message, response_message);
+
+            // We receive an incoming close message
+            control
+                .incoming
+                .send(Ok(Message::Close(None)))
+                .await
+                .unwrap();
+            // Then we expect the outgoing to close with no message
+            assert!(control.outgoing.recv().await.is_none());
+            assert!(control.incoming.is_closed());
+        })
+        .await;
+    }
+
+    #[tokio::test]
+    async fn callback_registered() {
+        let _ = env_logger::builder()
+            .is_test(true)
+            .filter_level(LevelFilter::Trace)
+            .try_init();
+
+        assert_client_interaction(|sender, mut control, _| async move {
+            let (callback_tx, mut callback_rx) = mpsc::channel(1);
+            let msg = Uuid::new_v4();
+            sender
+                .send(OutgoingMessage {
+                    msg: RequestMessage {
+                        uuid: msg,
+                        response: None,
+                        payload: TelemetryDefinitionRequest {
+                            name: "".to_string(),
+                            data_type: DataType::Float32,
+                        }
+                        .into(),
+                    },
+                    callback: Callback::Registered(callback_tx),
+                })
+                .await
+                .unwrap();
+
+            // We expect an outgoing message
+            assert!(matches!(
+                control.outgoing.recv().await.unwrap(),
+                Message::Text(_)
+            ));
+
+            // We handle the callback a few times
+            for _ in 0..5 {
+                // Then we get an incoming message for this callback
+                let response_message = ResponseMessage {
+                    uuid: Uuid::new_v4(),
+                    response: Some(msg),
+                    payload: TelemetryDefinitionResponse {
+                        uuid: Uuid::new_v4(),
+                    }
+                    .into(),
+                };
+                control
+                    .incoming
+                    .send(Ok(Message::Text(
+                        serde_json::to_string(&response_message).unwrap().into(),
+                    )))
+                    .await
+                    .unwrap();
+
+                // We expect the response
+                let (rx, responder) = callback_rx.recv().await.unwrap();
+                // And give us the message we provided it
+                assert_eq!(rx, response_message);
+                // Then the response gets sent out
+                responder
+                    .send(
+                        TelemetryDefinitionRequest {
+                            name: "".to_string(),
+                            data_type: DataType::Float32,
+                        }
+                        .into(),
+                    )
+                    .unwrap();
+
+                // We expect an outgoing message
+                assert!(matches!(
+                    control.outgoing.recv().await.unwrap(),
+                    Message::Text(_)
+                ));
+            }
+
+            // We receive an incoming close message
+            control
+                .incoming
+                .send(Ok(Message::Close(None)))
+                .await
+                .unwrap();
+            // Then we expect the outgoing to close with no message
+            assert!(control.outgoing.recv().await.is_none());
+            assert!(control.incoming.is_closed());
+        })
+        .await;
+    }
+
+    #[tokio::test]
+    async fn ping_pong() {
+        let _ = env_logger::builder()
+            .is_test(true)
+            .filter_level(LevelFilter::Trace)
+            .try_init();
+
+        assert_client_interaction(|_, mut control, _| async move {
+            // Expect a pong in response to a ping
+            let bytes = Bytes::from_owner(Uuid::new_v4().into_bytes());
+            control
+                .incoming
+                .send(Ok(Message::Ping(bytes.clone())))
+                .await
+                .unwrap();
+            let Some(Message::Pong(pong_bytes)) = control.outgoing.recv().await else {
+                panic!("Expected Pong Response");
+            };
+            assert_eq!(bytes, pong_bytes);
+
+            // Nothing should happen
+            control
+                .incoming
+                .send(Ok(Message::Pong(bytes.clone())))
+                .await
+                .unwrap();
+
+            // We receive an incoming close message
+            control
+                .incoming
+                .send(Ok(Message::Close(None)))
+                .await
+                .unwrap();
+            // Then we expect the outgoing to close with no message
+            assert!(control.outgoing.recv().await.is_none());
+            assert!(control.incoming.is_closed());
+        })
+        .await;
+    }
+}

api/src/client/error.rs

@@ -0,0 +1,37 @@
+use api_core::data_type::DataType;
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum ConnectError {
+    #[error(transparent)]
+    TungsteniteError(#[from] tokio_tungstenite::tungstenite::Error),
+    #[error(transparent)]
+    IoError(#[from] std::io::Error),
+}
+
+#[derive(Error, Debug)]
+pub enum MessageError {
+    #[error(transparent)]
+    TokioSendError(#[from] tokio::sync::mpsc::error::SendError<()>),
+    #[error(transparent)]
+    TokioTrySendError(#[from] tokio::sync::mpsc::error::TrySendError<()>),
+    #[error(transparent)]
+    TokioLockError(#[from] tokio::sync::TryLockError),
+    #[error("Incorrect Data Type. {expected} expected. {actual} actual.")]
+    IncorrectDataType {
+        expected: DataType,
+        actual: DataType,
+    },
+}
+
+#[derive(Error, Debug)]
+pub enum RequestError<E> {
+    #[error(transparent)]
+    TokioSendError(#[from] tokio::sync::mpsc::error::SendError<()>),
+    #[error(transparent)]
+    TokioLockError(#[from] tokio::sync::TryLockError),
+    #[error(transparent)]
+    RecvError(#[from] tokio::sync::oneshot::error::RecvError),
+    #[error(transparent)]
+    Inner(E),
+}

api/src/client/mod.rs

@@ -0,0 +1,598 @@
+pub mod command;
+mod config;
+mod context;
+pub mod error;
+pub mod telemetry;
+
+use crate::client::config::ClientConfiguration;
+use crate::client::error::{MessageError, RequestError};
+use crate::messages::callback::GenericCallbackError;
+use crate::messages::payload::RequestMessagePayload;
+use crate::messages::payload::ResponseMessagePayload;
+use crate::messages::{
+    ClientMessage, RegisterCallback, RequestMessage, RequestResponse, ResponseMessage,
+};
+use context::ClientContext;
+use error::ConnectError;
+use std::sync::Arc;
+use tokio::spawn;
+use tokio::sync::{mpsc, oneshot, watch, RwLock};
+use tokio_tungstenite::tungstenite::client::IntoClientRequest;
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+type RegisteredCallback = mpsc::Sender<(ResponseMessage, oneshot::Sender<RequestMessagePayload>)>;
+type ClientChannel = Arc<RwLock<mpsc::Sender<OutgoingMessage>>>;
+
+#[derive(Debug)]
+enum Callback {
+    None,
+    Once(oneshot::Sender<ResponseMessage>),
+    Registered(RegisteredCallback),
+}
+
+#[derive(Debug)]
+struct OutgoingMessage {
+    msg: RequestMessage,
+    callback: Callback,
+}
+
+pub struct Client {
+    cancel: CancellationToken,
+    channel: ClientChannel,
+    connected_state_rx: watch::Receiver<bool>,
+}
+
+impl Client {
+    pub fn connect<R>(request: R) -> Result<Self, ConnectError>
+    where
+        R: IntoClientRequest,
+    {
+        Self::connect_with_config(request, ClientConfiguration::default())
+    }
+
+    pub fn connect_with_config<R>(
+        request: R,
+        config: ClientConfiguration,
+    ) -> Result<Self, ConnectError>
+    where
+        R: IntoClientRequest,
+    {
+        let (tx, _rx) = mpsc::channel(1);
+        let cancel = CancellationToken::new();
+        let channel = Arc::new(RwLock::new(tx));
+        let (connected_state_tx, connected_state_rx) = watch::channel(false);
+        let context = ClientContext {
+            cancel: cancel.clone(),
+            request: request.into_client_request()?,
+            connected_state_tx,
+            client_configuration: config,
+        };
+
+        context.start(channel.clone())?;
+
+        Ok(Self {
+            cancel,
+            channel,
+            connected_state_rx,
+        })
+    }
+
+    pub async fn send_message<M: ClientMessage>(&self, msg: M) -> Result<(), MessageError> {
+        let sender = self.channel.read().await;
+        let data = sender.reserve().await?;
+        data.send(OutgoingMessage {
+            msg: RequestMessage {
+                uuid: Uuid::new_v4(),
+                response: None,
+                payload: msg.into(),
+            },
+            callback: Callback::None,
+        });
+        Ok(())
+    }
+
+    pub async fn send_message_if_connected<M: ClientMessage>(
+        &self,
+        msg: M,
+    ) -> Result<(), MessageError> {
+        let sender = self.channel.try_read()?;
+        let data = sender.reserve().await?;
+        data.send(OutgoingMessage {
+            msg: RequestMessage {
+                uuid: Uuid::new_v4(),
+                response: None,
+                payload: msg.into(),
+            },
+            callback: Callback::None,
+        });
+        Ok(())
+    }
+
+    pub fn try_send_message<M: ClientMessage>(&self, msg: M) -> Result<(), MessageError> {
+        let sender = self.channel.try_read()?;
+        let data = sender.try_reserve()?;
+        data.send(OutgoingMessage {
+            msg: RequestMessage {
+                uuid: Uuid::new_v4(),
+                response: None,
+                payload: msg.into(),
+            },
+            callback: Callback::None,
+        });
+        Ok(())
+    }
+
+    pub async fn send_request<M: RequestResponse>(
+        &self,
+        msg: M,
+    ) -> Result<M::Response, RequestError<<M::Response as TryFrom<ResponseMessagePayload>>::Error>>
+    {
+        let sender = self.channel.read().await;
+        let data = sender.reserve().await?;
+
+        let (tx, rx) = oneshot::channel();
+        data.send(OutgoingMessage {
+            msg: RequestMessage {
+                uuid: Uuid::new_v4(),
+                response: None,
+                payload: msg.into(),
+            },
+            callback: Callback::Once(tx),
+        });
+
+        let response = rx.await?;
+        let response = M::Response::try_from(response.payload).map_err(RequestError::Inner)?;
+
+        Ok(response)
+    }
+
+    pub async fn register_callback_channel<M: RegisterCallback>(
+        &self,
+        msg: M,
+    ) -> Result<mpsc::Receiver<(M::Callback, oneshot::Sender<M::Response>)>, MessageError>
+    where
+        <M as RegisterCallback>::Callback: Send + 'static,
+        <M as RegisterCallback>::Response: Send + 'static,
+        <<M as RegisterCallback>::Callback as TryFrom<ResponseMessagePayload>>::Error: Send,
+    {
+        let sender = self.channel.read().await;
+        let data = sender.reserve().await?;
+
+        let (inner_tx, mut inner_rx) = mpsc::channel(16);
+        let (outer_tx, outer_rx) = mpsc::channel(1);
+
+        data.send(OutgoingMessage {
+            msg: RequestMessage {
+                uuid: Uuid::new_v4(),
+                response: None,
+                payload: msg.into(),
+            },
+            callback: Callback::Registered(inner_tx),
+        });
+
+        spawn(async move {
+            // If the handler was unregistered we can stop
+            while let Some((msg, responder)) = inner_rx.recv().await {
+                let response: RequestMessagePayload = match M::Callback::try_from(msg.payload) {
+                    Err(_) => GenericCallbackError::MismatchedType.into(),
+                    Ok(o) => {
+                        let (response_tx, response_rx) = oneshot::channel::<M::Response>();
+                        match outer_tx.send((o, response_tx)).await {
+                            Err(_) => GenericCallbackError::CallbackClosed.into(),
+                            Ok(()) => response_rx
+                                .await
+                                .map(M::Response::into)
+                                .unwrap_or_else(|_| GenericCallbackError::CallbackClosed.into()),
+                        }
+                    }
+                };
+
+                if responder.send(response).is_err() {
+                    // If the callback was unregistered we can stop
+                    break;
+                }
+            }
+            println!("Exited Loop");
+        });
+
+        Ok(outer_rx)
+    }
+
+    pub async fn register_callback_fn<M: RegisterCallback, F>(
+        &self,
+        msg: M,
+        mut f: F,
+    ) -> Result<(), MessageError>
+    where
+        F: FnMut(M::Callback) -> M::Response + Send + 'static,
+    {
+        let sender = self.channel.read().await;
+        let data = sender.reserve().await?;
+
+        let (inner_tx, mut inner_rx) = mpsc::channel(16);
+
+        data.send(OutgoingMessage {
+            msg: RequestMessage {
+                uuid: Uuid::new_v4(),
+                response: None,
+                payload: msg.into(),
+            },
+            callback: Callback::Registered(inner_tx),
+        });
+
+        spawn(async move {
+            // If the handler was unregistered we can stop
+            while let Some((msg, responder)) = inner_rx.recv().await {
+                let response: RequestMessagePayload = match M::Callback::try_from(msg.payload) {
+                    Err(_) => GenericCallbackError::MismatchedType.into(),
+                    Ok(o) => f(o).into(),
+                };
+
+                if responder.send(response).is_err() {
+                    // If the callback was unregistered we can stop
+                    break;
+                }
+            }
+        });
+
+        Ok(())
+    }
+
+    pub async fn wait_connected(&self) {
+        let mut connected_rx = self.connected_state_rx.clone();
+
+        // If we aren't currently connected
+        if !*connected_rx.borrow_and_update() {
+            // Wait for a change notification
+            // If the channel is closed there is nothing we can do
+            let _ = connected_rx.changed().await;
+        }
+    }
+
+    pub async fn wait_disconnected(&self) {
+        let mut connected_rx = self.connected_state_rx.clone();
+
+        // If we are currently connected
+        if *connected_rx.borrow_and_update() {
+            // Wait for a change notification
+            // If the channel is closed there is nothing we can do
+            let _ = connected_rx.changed().await;
+        }
+    }
+}
+
+impl Drop for Client {
+    fn drop(&mut self) {
+        self.cancel.cancel();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::messages::command::CommandResponse;
+    use crate::messages::telemetry_definition::{
+        TelemetryDefinitionRequest, TelemetryDefinitionResponse,
+    };
+    use crate::messages::telemetry_entry::TelemetryEntry;
+    use api_core::command::{Command, CommandDefinition, CommandHeader};
+    use api_core::data_type::DataType;
+    use chrono::Utc;
+    use futures_util::future::{select, Either};
+    use futures_util::FutureExt;
+    use std::pin::pin;
+    use std::time::Duration;
+    use tokio::join;
+    use tokio::time::{sleep, timeout};
+
+    pub fn create_test_client() -> (mpsc::Receiver<OutgoingMessage>, watch::Sender<bool>, Client) {
+        let cancel = CancellationToken::new();
+        let (tx, rx) = mpsc::channel(1);
+        let channel = Arc::new(RwLock::new(tx));
+        let (connected_state_tx, connected_state_rx) = watch::channel(true);
+        let client = Client {
+            cancel,
+            channel,
+            connected_state_rx,
+        };
+        (rx, connected_state_tx, client)
+    }
+
+    #[tokio::test]
+    async fn send_message() {
+        let (mut rx, _, client) = create_test_client();
+
+        let msg_to_send = TelemetryEntry {
+            uuid: Uuid::new_v4(),
+            value: 0.0f32.into(),
+            timestamp: Utc::now(),
+        };
+        let msg_send = timeout(
+            Duration::from_secs(1),
+            client.send_message(msg_to_send.clone()),
+        );
+        let msg_recv = timeout(Duration::from_secs(1), rx.recv());
+
+        let (send, recv) = join!(msg_send, msg_recv);
+        send.unwrap().unwrap();
+        let recv = recv.unwrap().unwrap();
+
+        assert!(matches!(recv.callback, Callback::None));
+        assert!(recv.msg.response.is_none());
+        // uuid should be random
+
+        let RequestMessagePayload::TelemetryEntry(recv) = recv.msg.payload else {
+            panic!("Wrong Message Received")
+        };
+
+        assert_eq!(recv, msg_to_send);
+    }
+
+    #[tokio::test]
+    async fn send_message_if_connected() {
+        let (mut rx, _, client) = create_test_client();
+
+        let msg_to_send = TelemetryEntry {
+            uuid: Uuid::new_v4(),
+            value: 0.0f32.into(),
+            timestamp: Utc::now(),
+        };
+        let msg_send = timeout(
+            Duration::from_secs(1),
+            client.send_message_if_connected(msg_to_send.clone()),
+        );
+        let msg_recv = timeout(Duration::from_secs(1), rx.recv());
+
+        let (send, recv) = join!(msg_send, msg_recv);
+        send.unwrap().unwrap();
+        let recv = recv.unwrap().unwrap();
+
+        assert!(matches!(recv.callback, Callback::None));
+        assert!(recv.msg.response.is_none());
+        // uuid should be random
+
+        let RequestMessagePayload::TelemetryEntry(recv) = recv.msg.payload else {
+            panic!("Wrong Message Received")
+        };
+
+        assert_eq!(recv, msg_to_send);
+    }
+
+    #[tokio::test]
+    async fn send_message_if_connected_not_connected() {
+        let (_, connected_state_tx, client) = create_test_client();
+
+        let _lock = client.channel.write().await;
+        connected_state_tx.send_replace(false);
+
+        let msg_to_send = TelemetryEntry {
+            uuid: Uuid::new_v4(),
+            value: 0.0f32.into(),
+            timestamp: Utc::now(),
+        };
+        let msg_send = timeout(
+            Duration::from_secs(1),
+            client.send_message_if_connected(msg_to_send.clone()),
+        );
+
+        let Err(MessageError::TokioLockError(_)) = msg_send.await.unwrap() else {
+            panic!("Expected to Err due to lock being unavailable")
+        };
+    }
+
+    #[tokio::test]
+    async fn try_send_message() {
+        let (_tx, _, client) = create_test_client();
+
+        let msg_to_send = TelemetryEntry {
+            uuid: Uuid::new_v4(),
+            value: 0.0f32.into(),
+            timestamp: Utc::now(),
+        };
+        client.try_send_message(msg_to_send.clone()).unwrap();
+        let Err(MessageError::TokioTrySendError(_)) = client.try_send_message(msg_to_send.clone())
+        else {
+            panic!("Expected the buffer to be full");
+        };
+    }
+
+    #[tokio::test]
+    async fn send_request() {
+        let (mut tx, _, client) = create_test_client();
+
+        let msg_to_send = TelemetryDefinitionRequest {
+            name: "".to_string(),
+            data_type: DataType::Float32,
+        };
+        let response = timeout(
+            Duration::from_secs(1),
+            client.send_request(msg_to_send.clone()),
+        );
+
+        let response_uuid = Uuid::new_v4();
+        let outgoing_rx = timeout(Duration::from_secs(1), async {
+            let msg = tx.recv().await.unwrap();
+            let Callback::Once(cb) = msg.callback else {
+                panic!("Wrong Callback Type")
+            };
+            cb.send(ResponseMessage {
+                uuid: Uuid::new_v4(),
+                response: Some(msg.msg.uuid),
+                payload: TelemetryDefinitionResponse {
+                    uuid: response_uuid,
+                }
+                .into(),
+            })
+            .unwrap();
+        });
+
+        let (response, outgoing_rx) = join!(response, outgoing_rx);
+        let response = response.unwrap().unwrap();
+        outgoing_rx.unwrap();
+
+        assert_eq!(response.uuid, response_uuid);
+    }
+
+    #[tokio::test]
+    async fn register_callback_channel() {
+        let (mut tx, _, client) = create_test_client();
+
+        let msg_to_send = CommandDefinition {
+            name: "".to_string(),
+            parameters: vec![],
+        };
+        let mut response = timeout(
+            Duration::from_secs(1),
+            client.register_callback_channel(msg_to_send),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+
+        let outgoing_rx = timeout(Duration::from_secs(1), async {
+            let msg = tx.recv().await.unwrap();
+            let Callback::Registered(cb) = msg.callback else {
+                panic!("Wrong Callback Type")
+            };
+
+            // Check that we get responses to the callback the expected number of times
+            for i in 0..5 {
+                let (tx, rx) = oneshot::channel();
+                cb.send((
+                    ResponseMessage {
+                        uuid: Uuid::new_v4(),
+                        response: Some(msg.msg.uuid),
+                        payload: Command {
+                            header: CommandHeader {
+                                timestamp: Utc::now(),
+                            },
+                            parameters: Default::default(),
+                        }
+                        .into(),
+                    },
+                    tx,
+                ))
+                .await
+                .unwrap();
+                let RequestMessagePayload::CommandResponse(response) = rx.await.unwrap() else {
+                    panic!("Unexpected Response Type");
+                };
+                assert_eq!(response.response, format!("{i}"));
+            }
+        });
+
+        let responder = timeout(Duration::from_secs(1), async {
+            for i in 0..5 {
+                let (_cmd, responder) = response.recv().await.unwrap();
+                responder
+                    .send(CommandResponse {
+                        success: false,
+                        response: format!("{i}"),
+                    })
+                    .unwrap();
+            }
+        });
+
+        let (response, outgoing_rx) = join!(responder, outgoing_rx);
+        response.unwrap();
+        outgoing_rx.unwrap();
+    }
+
+    #[tokio::test]
+    async fn register_callback_fn() {
+        let (mut tx, _, client) = create_test_client();
+
+        let msg_to_send = CommandDefinition {
+            name: "".to_string(),
+            parameters: vec![],
+        };
+        let mut index = 0usize;
+        timeout(
+            Duration::from_secs(1),
+            client.register_callback_fn(msg_to_send, move |_| {
+                index += 1;
+                CommandResponse {
+                    success: false,
+                    response: format!("{}", index - 1),
+                }
+            }),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+
+        timeout(Duration::from_secs(1), async {
+            let msg = tx.recv().await.unwrap();
+            let Callback::Registered(cb) = msg.callback else {
+                panic!("Wrong Callback Type")
+            };
+
+            // Check that we get responses to the callback the expected number of times
+            for i in 0..3 {
+                let (tx, rx) = oneshot::channel();
+                cb.send((
+                    ResponseMessage {
+                        uuid: Uuid::new_v4(),
+                        response: Some(msg.msg.uuid),
+                        payload: Command {
+                            header: CommandHeader {
+                                timestamp: Utc::now(),
+                            },
+                            parameters: Default::default(),
+                        }
+                        .into(),
+                    },
+                    tx,
+                ))
+                .await
+                .unwrap();
+                let RequestMessagePayload::CommandResponse(response) = rx.await.unwrap() else {
+                    panic!("Unexpected Response Type");
+                };
+                assert_eq!(response.response, format!("{i}"));
+            }
+        })
+        .await
+        .unwrap();
+    }
+
+    #[tokio::test]
+    async fn connected_disconnected() {
+        let (_, connected, client) = create_test_client();
+
+        // When we're connected we should return immediately
+        connected.send_replace(true);
+        client.wait_connected().now_or_never().unwrap();
+
+        // When we're disconnected we should return immediately
+        connected.send_replace(false);
+        client.wait_disconnected().now_or_never().unwrap();
+
+        let c2 = connected.clone();
+        // When we're disconnected, we should not return immediately
+        let f1 = pin!(client.wait_connected());
+        let f2 = pin!(async move {
+            sleep(Duration::from_millis(1)).await;
+            c2.send_replace(true);
+        });
+        let r = select(f1, f2).await;
+        match r {
+            Either::Left(_) => panic!("Wait Connected Finished Before Connection Changed"),
+            Either::Right((_, other)) => timeout(Duration::from_secs(1), other).await.unwrap(),
+        }
+
+        let c2 = connected.clone();
+        // When we're disconnected, we should not return immediately
+        let f1 = pin!(client.wait_disconnected());
+        let f2 = pin!(async move {
+            sleep(Duration::from_millis(1)).await;
+            c2.send_replace(false);
+        });
+        let r = select(f1, f2).await;
+        match r {
+            Either::Left(_) => panic!("Wait Disconnected Finished Before Connection Changed"),
+            Either::Right((_, other)) => timeout(Duration::from_secs(1), other).await.unwrap(),
+        }
+    }
+}

api/src/client/telemetry.rs

@@ -0,0 +1,464 @@
+use crate::client::error::MessageError;
+use crate::client::Client;
+use crate::data_value::DataValue;
+use crate::messages::telemetry_definition::TelemetryDefinitionRequest;
+use crate::messages::telemetry_entry::TelemetryEntry;
+use api_core::data_type::{DataType, ToDataType};
+use chrono::{DateTime, Utc};
+use std::marker::PhantomData;
+use std::sync::Arc;
+use tokio::sync::{oneshot, RwLock};
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+pub struct TelemetryRegistry {
+    client: Arc<Client>,
+}
+
+impl TelemetryRegistry {
+    pub fn new(client: Arc<Client>) -> Self {
+        Self { client }
+    }
+
+    #[inline]
+    pub async fn register_generic(
+        &self,
+        name: impl Into<String>,
+        data_type: DataType,
+    ) -> GenericTelemetryHandle {
+        // inner for compilation performance
+        async fn inner(
+            client: Arc<Client>,
+            name: String,
+            data_type: DataType,
+        ) -> GenericTelemetryHandle {
+            let cancellation_token = CancellationToken::new();
+            let cancel_token = cancellation_token.clone();
+            let stored_client = client.clone();
+
+            let response_uuid = Arc::new(RwLock::new(Uuid::nil()));
+
+            let response_uuid_inner = response_uuid.clone();
+            let (tx, rx) = oneshot::channel();
+
+            tokio::spawn(async move {
+                let mut write_lock = Some(response_uuid_inner.write().await);
+                let _ = tx.send(());
+                while !cancel_token.is_cancelled() {
+                    if let Ok(response) = client
+                        .send_request(TelemetryDefinitionRequest {
+                            name: name.clone(),
+                            data_type,
+                        })
+                        .await
+                    {
+                        let mut lock = match write_lock {
+                            None => response_uuid_inner.write().await,
+                            Some(lock) => lock,
+                        };
+                        // Update the value in the lock
+                        *lock = response.uuid;
+                        // Set this value so the loop works
+                        write_lock = None;
+                    }
+
+                    client.wait_disconnected().await;
+                }
+            });
+
+            // Wait until the write lock is acquired
+            let _ = rx.await;
+            // Wait until the write lock is released for the first time
+            drop(response_uuid.read().await);
+
+            GenericTelemetryHandle {
+                cancellation_token,
+                uuid: response_uuid,
+                client: stored_client,
+                data_type,
+            }
+        }
+        inner(self.client.clone(), name.into(), data_type).await
+    }
+
+    #[inline]
+    pub async fn register<T: ToDataType>(&self, name: impl Into<String>) -> TelemetryHandle<T> {
+        self.register_generic(name, T::DATA_TYPE).await.coerce()
+    }
+}
+
+impl Drop for GenericTelemetryHandle {
+    fn drop(&mut self) {
+        self.cancellation_token.cancel();
+    }
+}
+
+pub struct GenericTelemetryHandle {
+    cancellation_token: CancellationToken,
+    uuid: Arc<RwLock<Uuid>>,
+    client: Arc<Client>,
+    data_type: DataType,
+}
+
+impl GenericTelemetryHandle {
+    pub async fn publish(
+        &self,
+        value: DataValue,
+        timestamp: DateTime<Utc>,
+    ) -> Result<(), MessageError> {
+        if value.to_data_type() != self.data_type {
+            return Err(MessageError::IncorrectDataType {
+                expected: self.data_type,
+                actual: value.to_data_type(),
+            });
+        }
+        let Ok(lock) = self.uuid.try_read() else {
+            return Ok(());
+        };
+        let uuid = *lock;
+        drop(lock);
+
+        self.client
+            .send_message_if_connected(TelemetryEntry {
+                uuid,
+                value,
+                timestamp,
+            })
+            .await
+            .or_else(|e| match e {
+                MessageError::TokioLockError(_) => Ok(()),
+                e => Err(e),
+            })?;
+
+        Ok(())
+    }
+
+    #[inline]
+    pub async fn publish_now(&self, value: DataValue) -> Result<(), MessageError> {
+        self.publish(value, Utc::now()).await
+    }
+
+    fn coerce<T: Into<DataValue>>(self) -> TelemetryHandle<T> {
+        TelemetryHandle::<T> {
+            generic_handle: self,
+            _phantom: PhantomData,
+        }
+    }
+}
+
+pub struct TelemetryHandle<T> {
+    generic_handle: GenericTelemetryHandle,
+    _phantom: PhantomData<T>,
+}
+
+impl<T> TelemetryHandle<T> {
+    pub fn to_generic(self) -> GenericTelemetryHandle {
+        self.generic_handle
+    }
+    pub fn as_generic(&self) -> &GenericTelemetryHandle {
+        &self.generic_handle
+    }
+}
+
+impl<T: Into<DataValue>> TelemetryHandle<T> {
+    #[inline]
+    pub async fn publish(&self, value: T, timestamp: DateTime<Utc>) -> Result<(), MessageError> {
+        self.as_generic().publish(value.into(), timestamp).await
+    }
+
+    #[inline]
+    pub async fn publish_now(&self, value: T) -> Result<(), MessageError> {
+        self.publish(value, Utc::now()).await
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::client::error::MessageError;
+    use crate::client::telemetry::TelemetryRegistry;
+    use crate::client::tests::create_test_client;
+    use crate::client::Callback;
+    use crate::messages::payload::RequestMessagePayload;
+    use crate::messages::telemetry_definition::{
+        TelemetryDefinitionRequest, TelemetryDefinitionResponse,
+    };
+    use crate::messages::telemetry_entry::TelemetryEntry;
+    use crate::messages::ResponseMessage;
+    use api_core::data_type::DataType;
+    use api_core::data_value::DataValue;
+    use futures_util::FutureExt;
+    use std::sync::Arc;
+    use std::time::Duration;
+    use tokio::task::yield_now;
+    use tokio::time::timeout;
+    use tokio::try_join;
+    use uuid::Uuid;
+
+    #[tokio::test]
+    async fn generic() {
+        // if _c drops then we are disconnected
+        let (mut rx, _c, client) = create_test_client();
+
+        let tlm = TelemetryRegistry::new(Arc::new(client));
+        let tlm_handle = tlm.register_generic("generic", DataType::Float32);
+
+        let tlm_uuid = Uuid::new_v4();
+
+        let expected_rx = async {
+            let msg = rx.recv().await.unwrap();
+            let Callback::Once(responder) = msg.callback else {
+                panic!("Expected Once Callback");
+            };
+            assert!(msg.msg.response.is_none());
+            let RequestMessagePayload::TelemetryDefinitionRequest(TelemetryDefinitionRequest {
+                name,
+                data_type,
+            }) = msg.msg.payload
+            else {
+                panic!("Expected Telemetry Definition Request")
+            };
+            assert_eq!(name, "generic".to_string());
+            assert_eq!(data_type, DataType::Float32);
+            responder
+                .send(ResponseMessage {
+                    uuid: Uuid::new_v4(),
+                    response: Some(msg.msg.uuid),
+                    payload: TelemetryDefinitionResponse { uuid: tlm_uuid }.into(),
+                })
+                .unwrap();
+        };
+
+        let (tlm_handle, _) = try_join!(
+            timeout(Duration::from_secs(1), tlm_handle),
+            timeout(Duration::from_secs(1), expected_rx),
+        )
+        .unwrap();
+
+        assert_eq!(*tlm_handle.uuid.try_read().unwrap(), tlm_uuid);
+
+        // This should NOT block if there is space in the queue
+        tlm_handle
+            .publish_now(0.0f32.into())
+            .now_or_never()
+            .unwrap()
+            .unwrap();
+
+        let tlm_msg = timeout(Duration::from_secs(1), rx.recv())
+            .await
+            .unwrap()
+            .unwrap();
+        assert!(matches!(tlm_msg.callback, Callback::None));
+        match tlm_msg.msg.payload {
+            RequestMessagePayload::TelemetryEntry(TelemetryEntry { uuid, value, .. }) => {
+                assert_eq!(uuid, tlm_uuid);
+                assert_eq!(value, DataValue::Float32(0.0f32));
+            }
+            _ => panic!("Expected Telemetry Entry"),
+        }
+    }
+
+    #[tokio::test]
+    async fn mismatched_type() {
+        let (mut rx, _, client) = create_test_client();
+
+        let tlm = TelemetryRegistry::new(Arc::new(client));
+        let tlm_handle = tlm.register_generic("generic", DataType::Float32);
+
+        let tlm_uuid = Uuid::new_v4();
+
+        let expected_rx = async {
+            let msg = rx.recv().await.unwrap();
+            let Callback::Once(responder) = msg.callback else {
+                panic!("Expected Once Callback");
+            };
+            assert!(msg.msg.response.is_none());
+            let RequestMessagePayload::TelemetryDefinitionRequest(TelemetryDefinitionRequest {
+                name,
+                data_type,
+            }) = msg.msg.payload
+            else {
+                panic!("Expected Telemetry Definition Request")
+            };
+            assert_eq!(name, "generic".to_string());
+            assert_eq!(data_type, DataType::Float32);
+            responder
+                .send(ResponseMessage {
+                    uuid: Uuid::new_v4(),
+                    response: Some(msg.msg.uuid),
+                    payload: TelemetryDefinitionResponse { uuid: tlm_uuid }.into(),
+                })
+                .unwrap();
+        };
+
+        let (tlm_handle, _) = try_join!(
+            timeout(Duration::from_secs(1), tlm_handle),
+            timeout(Duration::from_secs(1), expected_rx),
+        )
+        .unwrap();
+
+        assert_eq!(*tlm_handle.uuid.try_read().unwrap(), tlm_uuid);
+
+        match timeout(
+            Duration::from_secs(1),
+            tlm_handle.publish_now(0.0f64.into()),
+        )
+        .await
+        .unwrap()
+        {
+            Err(MessageError::IncorrectDataType { expected, actual }) => {
+                assert_eq!(expected, DataType::Float32);
+                assert_eq!(actual, DataType::Float64);
+            }
+            _ => panic!("Error Expected"),
+        }
+    }
+
+    #[tokio::test]
+    async fn typed() {
+        // if _c drops then we are disconnected
+        let (mut rx, _c, client) = create_test_client();
+
+        let tlm = TelemetryRegistry::new(Arc::new(client));
+        let tlm_handle = tlm.register::<bool>("typed");
+
+        let tlm_uuid = Uuid::new_v4();
+
+        let expected_rx = async {
+            let msg = rx.recv().await.unwrap();
+            let Callback::Once(responder) = msg.callback else {
+                panic!("Expected Once Callback");
+            };
+            assert!(msg.msg.response.is_none());
+            let RequestMessagePayload::TelemetryDefinitionRequest(TelemetryDefinitionRequest {
+                name,
+                data_type,
+            }) = msg.msg.payload
+            else {
+                panic!("Expected Telemetry Definition Request")
+            };
+            assert_eq!(name, "typed".to_string());
+            assert_eq!(data_type, DataType::Boolean);
+            responder
+                .send(ResponseMessage {
+                    uuid: Uuid::new_v4(),
+                    response: Some(msg.msg.uuid),
+                    payload: TelemetryDefinitionResponse { uuid: tlm_uuid }.into(),
+                })
+                .unwrap();
+        };
+
+        let (tlm_handle, _) = try_join!(
+            timeout(Duration::from_secs(1), tlm_handle),
+            timeout(Duration::from_secs(1), expected_rx),
+        )
+        .unwrap();
+
+        assert_eq!(*tlm_handle.as_generic().uuid.try_read().unwrap(), tlm_uuid);
+
+        // This should NOT block if there is space in the queue
+        tlm_handle
+            .publish_now(true)
+            .now_or_never()
+            .unwrap()
+            .unwrap();
+        // This should block as there should not be space in the queue
+        assert!(tlm_handle
+            .publish_now(false)
+            .now_or_never()
+            .is_none());
+
+        let tlm_msg = timeout(Duration::from_secs(1), rx.recv())
+            .await
+            .unwrap()
+            .unwrap();
+        assert!(matches!(tlm_msg.callback, Callback::None));
+        match tlm_msg.msg.payload {
+            RequestMessagePayload::TelemetryEntry(TelemetryEntry { uuid, value, .. }) => {
+                assert_eq!(uuid, tlm_uuid);
+                assert_eq!(value, DataValue::Boolean(true));
+            }
+            _ => panic!("Expected Telemetry Entry"),
+        }
+
+        let _make_generic_again = tlm_handle.to_generic();
+    }
+
+    #[tokio::test]
+    async fn reconnect() {
+        // if _c drops then we are disconnected
+        let (mut rx, connected, client) = create_test_client();
+
+        let tlm = TelemetryRegistry::new(Arc::new(client));
+        let tlm_handle = tlm.register_generic("generic", DataType::Float32);
+
+        let tlm_uuid = Uuid::new_v4();
+
+        let expected_rx = async {
+            let msg = rx.recv().await.unwrap();
+            let Callback::Once(responder) = msg.callback else {
+                panic!("Expected Once Callback");
+            };
+            assert!(msg.msg.response.is_none());
+            let RequestMessagePayload::TelemetryDefinitionRequest(TelemetryDefinitionRequest {
+                name,
+                data_type,
+            }) = msg.msg.payload
+            else {
+                panic!("Expected Telemetry Definition Request")
+            };
+            assert_eq!(name, "generic".to_string());
+            assert_eq!(data_type, DataType::Float32);
+            responder
+                .send(ResponseMessage {
+                    uuid: Uuid::new_v4(),
+                    response: Some(msg.msg.uuid),
+                    payload: TelemetryDefinitionResponse { uuid: tlm_uuid }.into(),
+                })
+                .unwrap();
+        };
+
+        let (tlm_handle, _) = try_join!(
+            timeout(Duration::from_secs(1), tlm_handle),
+            timeout(Duration::from_secs(1), expected_rx),
+        )
+        .unwrap();
+
+        assert_eq!(*tlm_handle.uuid.try_read().unwrap(), tlm_uuid);
+
+        // Notify Disconnect
+        connected.send_replace(false);
+        // Notify Reconnect
+        connected.send_replace(true);
+
+        {
+            let new_tlm_uuid = Uuid::new_v4();
+
+            let msg = rx.recv().await.unwrap();
+            let Callback::Once(responder) = msg.callback else {
+                panic!("Expected Once Callback");
+            };
+            assert!(msg.msg.response.is_none());
+            let RequestMessagePayload::TelemetryDefinitionRequest(TelemetryDefinitionRequest {
+                name,
+                data_type,
+            }) = msg.msg.payload
+            else {
+                panic!("Expected Telemetry Definition Request")
+            };
+            assert_eq!(name, "generic".to_string());
+            assert_eq!(data_type, DataType::Float32);
+            responder
+                .send(ResponseMessage {
+                    uuid: Uuid::new_v4(),
+                    response: Some(msg.msg.uuid),
+                    payload: TelemetryDefinitionResponse { uuid: new_tlm_uuid }.into(),
+                })
+                .unwrap();
+
+            // Yield to the executor so that the UUIDs can be updated
+            yield_now().await;
+
+            assert_eq!(*tlm_handle.uuid.try_read().unwrap(), new_tlm_uuid);
+        }
+    }
+}

api/src/lib.rs

@@ -0,0 +1,15 @@
+pub mod client;
+pub mod data_type {
+    pub use api_core::data_type::*;
+}
+pub mod data_value {
+    pub use api_core::data_value::*;
+}
+pub mod messages;
+
+pub mod macros {
+    pub use api_proc_macro::IntoCommandDefinition;
+}
+
+#[cfg(test)]
+pub mod test;

api/src/messages/callback.rs

@@ -0,0 +1,7 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub enum GenericCallbackError {
+    CallbackClosed,
+    MismatchedType,
+}

api/src/messages/command.rs

@@ -0,0 +1,15 @@
+use crate::messages::RegisterCallback;
+use serde::{Deserialize, Serialize};
+
+pub use api_core::command::*;
+
+impl RegisterCallback for CommandDefinition {
+    type Callback = Command;
+    type Response = CommandResponse;
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct CommandResponse {
+    pub success: bool,
+    pub response: String,
+}

api/src/messages/mod.rs

@@ -0,0 +1,40 @@
+pub mod callback;
+pub mod command;
+pub mod payload;
+pub mod telemetry_definition;
+pub mod telemetry_entry;
+
+use payload::{RequestMessagePayload, ResponseMessagePayload};
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RequestMessage {
+    pub uuid: Uuid,
+    #[serde(default)]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub response: Option<Uuid>,
+    #[serde(flatten)]
+    pub payload: RequestMessagePayload,
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct ResponseMessage {
+    pub uuid: Uuid,
+    #[serde(default)]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub response: Option<Uuid>,
+    #[serde(flatten)]
+    pub payload: ResponseMessagePayload,
+}
+
+pub trait ClientMessage: Into<RequestMessagePayload> {}
+
+pub trait RequestResponse: Into<RequestMessagePayload> {
+    type Response: TryFrom<ResponseMessagePayload>;
+}
+
+pub trait RegisterCallback: Into<RequestMessagePayload> {
+    type Callback: TryFrom<ResponseMessagePayload>;
+    type Response: Into<RequestMessagePayload>;
+}

api/src/messages/payload.rs

@@ -0,0 +1,23 @@
+use crate::messages::callback::GenericCallbackError;
+use crate::messages::command::{Command, CommandDefinition, CommandResponse};
+use crate::messages::telemetry_definition::{
+    TelemetryDefinitionRequest, TelemetryDefinitionResponse,
+};
+use crate::messages::telemetry_entry::TelemetryEntry;
+use derive_more::{From, TryInto};
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, From)]
+pub enum RequestMessagePayload {
+    TelemetryDefinitionRequest(TelemetryDefinitionRequest),
+    TelemetryEntry(TelemetryEntry),
+    GenericCallbackError(GenericCallbackError),
+    CommandDefinition(CommandDefinition),
+    CommandResponse(CommandResponse),
+}
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, From, TryInto)]
+pub enum ResponseMessagePayload {
+    TelemetryDefinitionResponse(TelemetryDefinitionResponse),
+    Command(Command),
+}

api/src/messages/telemetry_definition.rs

@@ -0,0 +1,19 @@
+use crate::data_type::DataType;
+use crate::messages::RequestResponse;
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct TelemetryDefinitionRequest {
+    pub name: String,
+    pub data_type: DataType,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct TelemetryDefinitionResponse {
+    pub uuid: Uuid,
+}
+
+impl RequestResponse for TelemetryDefinitionRequest {
+    type Response = TelemetryDefinitionResponse;
+}

api/src/messages/telemetry_entry.rs

@@ -0,0 +1,14 @@
+use crate::data_value::DataValue;
+use crate::messages::ClientMessage;
+use chrono::{DateTime, Utc};
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct TelemetryEntry {
+    pub uuid: Uuid,
+    pub value: DataValue,
+    pub timestamp: DateTime<Utc>,
+}
+
+impl ClientMessage for TelemetryEntry {}

api/src/test/mock_stream_sink.rs

@@ -0,0 +1,82 @@
+use futures_util::sink::{unfold, Unfold};
+use futures_util::{Sink, SinkExt, Stream};
+use std::fmt::Display;
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+use tokio::sync::mpsc;
+use tokio::sync::mpsc::error::SendError;
+use tokio::sync::mpsc::{Receiver, Sender};
+
+pub struct MockStreamSinkControl<T, R> {
+    pub incoming: Sender<T>,
+    pub outgoing: Receiver<R>,
+}
+
+pub struct MockStreamSink<T, U1, U2> {
+    stream_rx: Receiver<T>,
+    sink_tx: Pin<Box<Unfold<u32, U1, U2>>>,
+}
+
+impl<T, U1, U2> Stream for MockStreamSink<T, U1, U2>
+where
+    Self: Unpin,
+{
+    type Item = T;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.stream_rx.poll_recv(cx)
+    }
+}
+
+impl<T, R, U1, U2, E> Sink<R> for MockStreamSink<T, U1, U2>
+where
+    U1: FnMut(u32, R) -> U2,
+    U2: Future<Output = Result<u32, E>>,
+{
+    type Error = E;
+
+    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.sink_tx.poll_ready_unpin(cx)
+    }
+
+    fn start_send(mut self: Pin<&mut Self>, item: R) -> Result<(), Self::Error> {
+        self.sink_tx.start_send_unpin(item)
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.sink_tx.poll_flush_unpin(cx)
+    }
+
+    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.sink_tx.poll_close_unpin(cx)
+    }
+}
+
+pub fn create_mock_stream_sink<T: Send, R: Send + 'static>() -> (
+    MockStreamSinkControl<T, R>,
+    impl Stream<Item = T> + Sink<R, Error = impl Display>,
+) {
+    let (stream_tx, stream_rx) = mpsc::channel::<T>(1);
+    let (sink_tx, sink_rx) = mpsc::channel::<R>(1);
+
+    let sink_tx = Arc::new(sink_tx);
+
+    (
+        MockStreamSinkControl {
+            incoming: stream_tx,
+            outgoing: sink_rx,
+        },
+        MockStreamSink::<T, _, _> {
+            stream_rx,
+            sink_tx: Box::pin(unfold(0u32, move |_, item| {
+                let sink_tx = sink_tx.clone();
+                async move {
+                    sink_tx.send(item).await?;
+                    Ok(0u32) as Result<_, SendError<R>>
+                }
+            })),
+        },
+    )
+}

api/src/test/mod.rs

@@ -0,0 +1 @@
+pub mod mock_stream_sink;