test: Introduce crate flow-test

Cargo.toml

@@ -24,6 +24,7 @@ tokio = { version = "1.32.0", features = ["macros", "net", "rt", "sync"] }
 [workspace]
 resolver = "2"
 members = [
+    "flow-test",
     "proxy",
     "tag-generator",
     "tasks",

flow-test/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "flow-test"
+version = "0.1.0"
+edition = "2021"
+license = "MIT OR Apache-2.0"
+
+[dependencies]
+bstr = { version = "1.9.0", default-features = false }
+bytes = "1.5.0"
+imap-codec = "2.0.0"
+imap-flow = { path = ".." }
+imap-types = "2.0.0"
+tokio = { version = "1.32.0", features = ["macros", "net", "rt", "time"] }
+tracing = "0.1.37"
+tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }

flow-test/README.md

@@ -0,0 +1,3 @@
+# flow-test
+
+Test harness for writing lightweight unit tests for `imap-flow`.

flow-test/src/client_tester.rs

@@ -0,0 +1,155 @@
+use std::net::SocketAddr;
+
+use bstr::ByteSlice;
+use imap_flow::{
+    client::{ClientFlow, ClientFlowError, ClientFlowEvent, ClientFlowOptions},
+    stream::AnyStream,
+};
+use imap_types::bounded_static::ToBoundedStatic;
+use tokio::net::TcpStream;
+use tracing::trace;
+
+use crate::codecs::Codecs;
+
+/// A wrapper for `ClientFlow` suitable for testing.
+pub struct ClientTester {
+    codecs: Codecs,
+    client_flow_options: ClientFlowOptions,
+    connection_state: ConnectionState,
+}
+
+impl ClientTester {
+    pub async fn new(
+        codecs: Codecs,
+        client_flow_options: ClientFlowOptions,
+        server_address: SocketAddr,
+    ) -> Self {
+        let stream = TcpStream::connect(server_address).await.unwrap();
+        trace!(?server_address, "Client is connected");
+        Self {
+            codecs,
+            client_flow_options,
+            connection_state: ConnectionState::Connected { stream },
+        }
+    }
+
+    pub async fn receive_greeting(&mut self, expected_bytes: &[u8]) {
+        let expected_greeting = self.codecs.decode_greeting(expected_bytes);
+        match self.connection_state.take() {
+            ConnectionState::Connected { stream } => {
+                let stream = AnyStream::new(stream);
+                let (client, greeting) =
+                    ClientFlow::receive_greeting(stream, self.client_flow_options.clone())
+                        .await
+                        .unwrap();
+                assert_eq!(expected_greeting, greeting);
+                self.connection_state = ConnectionState::Greeted { client };
+            }
+            ConnectionState::Greeted { .. } => {
+                panic!("Client is already greeted");
+            }
+            ConnectionState::Disconnected => {
+                panic!("Client is already disconnected");
+            }
+        }
+    }
+
+    pub async fn send_command(&mut self, bytes: &[u8]) {
+        let enqueued_command = self.codecs.decode_command_normalized(bytes);
+        let client = self.connection_state.greeted();
+        let enqueued_handle = client.enqueue_command(enqueued_command.to_static());
+        let event = client.progress().await.unwrap();
+        match event {
+            ClientFlowEvent::CommandSent { handle, command } => {
+                assert_eq!(enqueued_handle, handle);
+                assert_eq!(enqueued_command, command);
+            }
+            event => {
+                panic!("Client emitted unexpected event: {event:?}");
+            }
+        }
+    }
+
+    pub async fn receive_data(&mut self, expected_bytes: &[u8]) {
+        let expected_data = self.codecs.decode_data(expected_bytes);
+        let client = self.connection_state.greeted();
+        match client.progress().await.unwrap() {
+            ClientFlowEvent::DataReceived { data } => {
+                assert_eq!(expected_data, data);
+            }
+            event => {
+                panic!("Client emitted unexpected event: {event:?}");
+            }
+        }
+    }
+
+    pub async fn receive_status(&mut self, expected_bytes: &[u8]) {
+        let expected_status = self.codecs.decode_status(expected_bytes);
+        let client = self.connection_state.greeted();
+        match client.progress().await.unwrap() {
+            ClientFlowEvent::StatusReceived { status } => {
+                assert_eq!(expected_status, status);
+            }
+            event => {
+                panic!("Client emitted unexpected event: {event:?}");
+            }
+        }
+    }
+
+    pub async fn receive_error_because_malformed_message(&mut self, expected_bytes: &[u8]) {
+        let error = match self.connection_state.take() {
+            ConnectionState::Connected { stream } => {
+                let stream = AnyStream::new(stream);
+                ClientFlow::receive_greeting(stream, self.client_flow_options.clone())
+                    .await
+                    .unwrap_err()
+            }
+            ConnectionState::Greeted { mut client } => {
+                let error = client.progress().await.unwrap_err();
+                self.connection_state = ConnectionState::Greeted { client };
+                error
+            }
+            ConnectionState::Disconnected => {
+                panic!("Client is already disconnected")
+            }
+        };
+        match error {
+            ClientFlowError::MalformedMessage { discarded_bytes } => {
+                assert_eq!(expected_bytes.as_bstr(), discarded_bytes.as_bstr());
+            }
+            error => {
+                panic!("Client emitted unexpected error: {error:?}");
+            }
+        }
+    }
+}
+
+/// The current state of the connection between client and server.
+#[allow(clippy::large_enum_variant)]
+enum ConnectionState {
+    /// The client has established a TCP connection to the server.
+    Connected { stream: TcpStream },
+    /// The client was greeted by the server.
+    Greeted { client: ClientFlow },
+    /// The TCP connection between client and server was dropped.
+    Disconnected,
+}
+
+impl ConnectionState {
+    /// Assumes that the client was already greeted by the server and returns the `ClientFlow`.
+    fn greeted(&mut self) -> &mut ClientFlow {
+        match self {
+            ConnectionState::Connected { .. } => {
+                panic!("Client is not greeted yet");
+            }
+            ConnectionState::Greeted { client } => client,
+            ConnectionState::Disconnected => {
+                panic!("Client is already disconnected");
+            }
+        }
+    }
+
+    fn take(&mut self) -> ConnectionState {
+        std::mem::replace(self, ConnectionState::Disconnected)
+    }
+}

flow-test/src/codecs.rs

@@ -0,0 +1,178 @@
+use bstr::ByteSlice;
+use imap_codec::{
+    decode::Decoder, encode::Encoder, AuthenticateDataCodec, CommandCodec, GreetingCodec,
+    IdleDoneCodec, ResponseCodec,
+};
+use imap_types::{
+    command::Command,
+    response::{Data, Greeting, Response, Status},
+};
+
+/// Contains all codecs from `imap-codec`.
+#[derive(Clone, Debug, Default, PartialEq)]
+#[non_exhaustive]
+pub struct Codecs {
+    pub greeting_codec: GreetingCodec,
+    pub command_codec: CommandCodec,
+    pub response_codec: ResponseCodec,
+    pub authenticate_data_codec: AuthenticateDataCodec,
+    pub idle_done_codec: IdleDoneCodec,
+}
+
+impl Codecs {
+    pub fn encode_greeting(&self, greeting: &Greeting) -> Vec<u8> {
+        self.greeting_codec.encode(greeting).dump()
+    }
+
+    pub fn encode_command(&self, command: &Command) -> Vec<u8> {
+        self.command_codec.encode(command).dump()
+    }
+
+    pub fn encode_response(&self, response: &Response) -> Vec<u8> {
+        self.response_codec.encode(response).dump()
+    }
+
+    pub fn encode_data(&self, data: &Data) -> Vec<u8> {
+        self.response_codec
+            .encode(&Response::Data(data.clone()))
+            .dump()
+    }
+
+    pub fn encode_status(&self, status: &Status) -> Vec<u8> {
+        self.response_codec
+            .encode(&Response::Status(status.clone()))
+            .dump()
+    }
+
+    pub fn decode_greeting<'a>(&self, bytes: &'a [u8]) -> Greeting<'a> {
+        match self.greeting_codec.decode(bytes) {
+            Ok((rem, greeting)) => {
+                if !rem.is_empty() {
+                    panic!(
+                        "Expected single greeting but there are remaining bytes {:?}",
+                        rem.as_bstr()
+                    )
+                }
+                greeting
+            }
+            Err(err) => {
+                panic!(
+                    "Got error {:?} when parsing greeting from bytes {:?}",
+                    err,
+                    bytes.as_bstr()
+                )
+            }
+        }
+    }
+
+    pub fn decode_command<'a>(&self, bytes: &'a [u8]) -> Command<'a> {
+        match self.command_codec.decode(bytes) {
+            Ok((rem, command)) => {
+                if !rem.is_empty() {
+                    panic!(
+                        "Expected single command but there are remaining bytes {:?}",
+                        rem.as_bstr()
+                    )
+                }
+                command
+            }
+            Err(err) => {
+                panic!(
+                    "Got error {:?} when parsing command from bytes {:?}",
+                    err,
+                    bytes.as_bstr()
+                )
+            }
+        }
+    }
+
+    pub fn decode_response<'a>(&self, bytes: &'a [u8]) -> Response<'a> {
+        match self.response_codec.decode(bytes) {
+            Ok((rem, response)) => {
+                if !rem.is_empty() {
+                    panic!(
+                        "Expected single response but there are remaining bytes {:?}",
+                        rem.as_bstr()
+                    )
+                }
+                response
+            }
+            Err(err) => {
+                panic!(
+                    "Got error {:?} when parsing response bytes {:?}",
+                    err,
+                    bytes.as_bstr()
+                )
+            }
+        }
+    }
+
+    pub fn decode_data<'a>(&self, bytes: &'a [u8]) -> Data<'a> {
+        let Response::Data(expected_data) = self.decode_response(bytes) else {
+            panic!("Got wrong response type when parsing data from {bytes:?}")
+        };
+        expected_data
+    }
+
+    pub fn decode_status<'a>(&self, bytes: &'a [u8]) -> Status<'a> {
+        let Response::Status(expected_status) = self.decode_response(bytes) else {
+            panic!("Got wrong response type when parsing status from {bytes:?}")
+        };
+        expected_status
+    }
+
+    pub fn decode_greeting_normalized<'a>(&self, bytes: &'a [u8]) -> Greeting<'a> {
+        let greeting = self.decode_greeting(bytes);
+        let normalized_bytes = self.encode_greeting(&greeting);
+        assert_eq!(
+            normalized_bytes.as_bstr(),
+            bytes.as_bstr(),
+            "Bytes must contain a normalized greeting"
+        );
+        greeting
+    }
+
+    pub fn decode_command_normalized<'a>(&self, bytes: &'a [u8]) -> Command<'a> {
+        let command = self.decode_command(bytes);
+        let normalized_bytes = self.encode_command(&command);
+        assert_eq!(
+            normalized_bytes.as_bstr(),
+            bytes.as_bstr(),
+            "Bytes must contain a normalized command"
+        );
+        command
+    }
+
+    pub fn decode_response_normalized<'a>(&self, bytes: &'a [u8]) -> Response<'a> {
+        let response = self.decode_response(bytes);
+        let normalized_bytes = self.encode_response(&response);
+        assert_eq!(
+            normalized_bytes.as_bstr(),
+            bytes.as_bstr(),
+            "Bytes must contain a normalized response"
+        );
+        response
+    }
+
+    pub fn decode_data_normalized<'a>(&self, bytes: &'a [u8]) -> Data<'a> {
+        let data = self.decode_data(bytes);
+        let normalized_bytes = self.encode_data(&data);
+        assert_eq!(
+            normalized_bytes.as_bstr(),
+            bytes.as_bstr(),
+            "Bytes must contain a normalized data"
+        );
+        data
+    }
+
+    pub fn decode_status_normalized<'a>(&self, bytes: &'a [u8]) -> Status<'a> {
+        let status = self.decode_status(bytes);
+        let normalized_bytes = self.encode_status(&status);
+        assert_eq!(
+            normalized_bytes.as_bstr(),
+            bytes.as_bstr(),
+            "Bytes must contain a normalized status"
+        );
+        status
+    }
+}

flow-test/src/lib.rs

@@ -0,0 +1,6 @@
+pub mod client_tester;
+pub mod codecs;
+pub mod mock;
+pub mod runtime;
+pub mod server_tester;
+pub mod test_setup;