Repository: quartiq/minimq
Pull request #146 (2023-08-28 10:57:28)

Refactoring builder to accept a single buffer
This PR simplifies usability of the `Config` structure builder by only requiring a single user buffer.

This is a follow-on to #145 to simplify builder arguments.


Patch:
diff --git a/src/broker.rs b/src/broker.rs
index 1d78526..9c5e32b 100644
--- a/src/broker.rs
+++ b/src/broker.rs
@@ -13,6 +13,7 @@ pub trait Broker {
 
 /// A broker that is specified using a qualified domain-name. The name will be resolved at some
 /// point in the future.
+#[derive(Debug)]
 pub struct NamedBroker<R: Dns, const T: usize = 253> {
     raw: heapless::String<T>,
     resolver: R,
@@ -63,6 +64,7 @@ impl<R: Dns, const T: usize> Broker for NamedBroker<R, T> {
 }
 
 /// A simple broker specification where the address of the broker is known in advance.
+#[derive(Debug, Copy, Clone, PartialEq)]
 pub struct IpBroker {
     addr: SocketAddr,
 }
diff --git a/src/config.rs b/src/config.rs
index acc8c2d..990a379 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -3,37 +3,57 @@ use core::convert::TryFrom;
 use embedded_time::duration::{Extensions, Milliseconds};
 use heapless::String;
 
-pub(crate) enum WillState<'a> {
-    BufferAvailable(&'a mut [u8]),
-    Serialized(SerializedWill<'a>),
+#[derive(Debug, Copy, Clone, PartialEq)]
+pub enum BufferConfig {
+    /// Specify a buffer as having a minimum size in bytes.
+    Minimum(usize),
+
+    /// Specify a buffer as having a maximum size in bytes.
+    Maximum(usize),
+
+    /// Specify a buffer as having exactly a size in bytes.
+    Exactly(usize),
 }
 
-/// Configuration specifying the operational state of the MQTT client.
-pub struct Config<'a, Broker: crate::Broker> {
+#[derive(Debug)]
+pub(crate) struct Config<'a, Broker: crate::Broker> {
     pub(crate) broker: Broker,
     pub(crate) rx_buffer: &'a mut [u8],
     pub(crate) tx_buffer: &'a mut [u8],
     pub(crate) state_buffer: &'a mut [u8],
-    pub(crate) will: Option<WillState<'a>>,
+    pub(crate) will: Option<SerializedWill<'a>>,
     pub(crate) client_id: String<64>,
     pub(crate) keepalive_interval: Milliseconds<u32>,
     pub(crate) downgrade_qos: bool,
 }
 
-impl<'a, Broker: crate::Broker> Config<'a, Broker> {
+/// Configuration specifying the operational state of the MQTT client.
+#[derive(Debug)]
+pub struct ConfigBuilder<'a, Broker: crate::Broker> {
+    buffer: &'a mut [u8],
+    broker: Broker,
+    rx_config: Option<BufferConfig>,
+    tx_config: Option<BufferConfig>,
+    session_state_config: Option<BufferConfig>,
+    will: Option<SerializedWill<'a>>,
+    client_id: String<64>,
+    keepalive_interval: Milliseconds<u32>,
+    downgrade_qos: bool,
+}
+
+impl<'a, Broker: crate::Broker> ConfigBuilder<'a, Broker> {
     /// Construct configuration for the MQTT client.
     ///
     /// # Args
-    /// * `rx` - Memory used for receiving messages. The length of this buffer is the maximum
-    /// receive packet length.
-    /// * `tx` - Memory used for transmitting messages. The length of this buffer is the max
-    /// transmit length.
-    pub fn new(broker: Broker, rx: &'a mut [u8], tx: &'a mut [u8]) -> Self {
+    /// * `buffer` - Memory used by the MQTT client. This memory is used for the will, the message
+    /// receive buffer, the transmission buffer, and the client session state.
+    pub fn new(broker: Broker, buffer: &'a mut [u8]) -> Self {
         Self {
             broker,
-            rx_buffer: rx,
-            tx_buffer: tx,
-            state_buffer: &mut [],
+            buffer,
+            session_state_config: None,
+            rx_config: None,
+            tx_config: None,
             client_id: String::new(),
             keepalive_interval: 59_000.milliseconds(),
             downgrade_qos: false,
@@ -41,15 +61,34 @@ impl<'a, Broker: crate::Broker> Config<'a, Broker> {
         }
     }
 
-    /// Provide additional buffer space if messages above [QoS::AtMostOnce] are required.
-    pub fn session_state(mut self, buffer: &'a mut [u8]) -> Self {
-        self.state_buffer = buffer;
+    /// Specify a specific configuration for the session state buffer.
+    ///
+    /// # Note
+    /// The session state buffer is used for publications greater than [QoS::AtMostOnce]. If these
+    /// messages are unused, you can specify [BufferConfig::Exactly(0)].
+    ///
+    /// # Args
+    /// * `config` - The configuration for the size of the session state buffer.
+    pub fn session_state(mut self, config: BufferConfig) -> Self {
+        self.session_state_config.replace(config);
+        self
+    }
+
+    /// Specify a specific configuration for the message receive buffer.
+    ///
+    /// # Args
+    /// * `config` - The configuration for the size of the receive buffer.
+    pub fn rx_buffer(mut self, config: BufferConfig) -> Self {
+        self.rx_config.replace(config);
         self
     }
 
-    /// Provide additional buffer space for an optional connection [Will]
-    pub fn will_buffer(mut self, buffer: &'a mut [u8]) -> Self {
-        self.will = Some(WillState::BufferAvailable(buffer));
+    /// Specify a specific configuration for the message transmit buffer.
+    ///
+    /// # Args
+    /// * `config` - The configuration for the size of the message transmit buffer.
+    pub fn tx_buffer(mut self, config: BufferConfig) -> Self {
+        self.tx_config.replace(config);
         self
     }
 
@@ -81,27 +120,168 @@ impl<'a, Broker: crate::Broker> Config<'a, Broker> {
         self
     }
 
-    /// Construct a will in the provided buffer.
-    pub fn will_with_buffer(
-        mut self,
-        buf: &'a mut [u8],
-        will: Will<'_>,
-    ) -> Result<Self, crate::ser::Error> {
-        self = self.will_buffer(buf);
-        self.will(will)
-    }
-
     /// Specify the Will message to be sent if the client disconnects.
     ///
     /// # Args
     /// * `will` - The will to use.
     pub fn will(mut self, will: Will<'_>) -> Result<Self, crate::ser::Error> {
-        let Some(WillState::BufferAvailable(buf)) = self.will else {
-            return Err(crate::ser::Error::InsufficientMemory);
+        let will_len = will.serialized_len();
+        let (head, tail) = self.buffer.split_at_mut(will_len);
+        self.buffer = tail;
+        self.will = Some(will.serialize(head)?);
+
+        Ok(self)
+    }
+
+    /// Consume the configuration and split the user-provided memory into the necessary buffers.
+    pub(crate) fn build(self) -> Config<'a, Broker> {
+        let configs = [self.rx_config, self.tx_config, self.session_state_config];
+        let mut buffers = [None, None, None];
+
+        let mut data = self.buffer;
+
+        // First, check for any exact configurations.
+        for (size, buffer) in configs
+            .iter()
+            .zip(buffers.iter_mut())
+            .filter_map(|(config, buf)| {
+                if let Some(BufferConfig::Exactly(size)) = config {
+                    Some((size, buf))
+                } else {
+                    None
+                }
+            })
+        {
+            let (head, tail) = data.split_at_mut(*size);
+            data = tail;
+            buffer.replace(head);
+        }
+
+        // Calculate the prescribed size for remaining buffers. This will be the base size used -
+        // any buffers that specify a min/max size will use this size and their respective mins or
+        // maxes.
+        let size = {
+            let remainder = buffers.iter().filter(|x| x.is_none()).count();
+            data.len().checked_div(remainder).unwrap_or(0)
         };
 
-        self.will = Some(WillState::Serialized(will.serialize(buf)?));
+        // Note: We intentionally are not dynamically updating the prescribed sizes as we allocate
+        // buffers because this would invalidate any previous potential min/max calculations that
+        // we conducted.
+        for (config, buffer) in configs
+            .iter()
+            .zip(buffers.iter_mut())
+            .filter_map(|(config, buf)| config.as_ref().map(|config| (config, buf)))
+        {
+            match config {
+                BufferConfig::Maximum(max) => {
+                    let (head, tail) = data.split_at_mut(size.min(*max));
+                    data = tail;
+                    buffer.replace(head);
+                }
+                BufferConfig::Minimum(min) => {
+                    let (head, tail) = data.split_at_mut(size.max(*min));
+                    data = tail;
+                    buffer.replace(head);
+                }
+                _ => {}
+            }
+        }
 
-        Ok(self)
+        // Any remaining buffers have no outstanding restrictions and the remaining memory will be
+        // evenly split amongst them.
+        let size = {
+            let remainder = buffers.iter().filter(|x| x.is_none()).count();
+            data.len().checked_div(remainder).unwrap_or(0)
+        };
+
+        for buffer in buffers.iter_mut().filter(|buf| buf.is_none()) {
+            let (head, tail) = data.split_at_mut(size);
+            data = tail;
+            buffer.replace(head);
+        }
+
+        Config {
+            broker: self.broker,
+            client_id: self.client_id,
+            will: self.will,
+            downgrade_qos: self.downgrade_qos,
+            keepalive_interval: self.keepalive_interval,
+            rx_buffer: buffers[0].take().unwrap(),
+            tx_buffer: buffers[1].take().unwrap(),
+            state_buffer: buffers[2].take().unwrap(),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::broker::IpBroker;
+
+    #[test]
+    pub fn basic_config() {
+        let mut buffer = [0; 30];
+        let localhost: embedded_nal::IpAddr = "127.0.0.1".parse().unwrap();
+        let builder: ConfigBuilder<IpBroker> = ConfigBuilder::new(localhost.into(), &mut buffer);
+
+        let config = builder.build();
+        assert!(config.tx_buffer.len() == 10);
+        assert!(config.rx_buffer.len() == 10);
+        assert!(config.state_buffer.len() == 10);
+    }
+
+    #[test]
+    pub fn without_session_state() {
+        let mut buffer = [0; 30];
+        let localhost: embedded_nal::IpAddr = "127.0.0.1".parse().unwrap();
+        let builder: ConfigBuilder<IpBroker> = ConfigBuilder::new(localhost.into(), &mut buffer)
+            .session_state(BufferConfig::Exactly(0));
+
+        let config = builder.build();
+        assert!(config.tx_buffer.len() == 15);
+        assert!(config.rx_buffer.len() == 15);
+        assert!(config.state_buffer.is_empty());
+    }
+
+    #[test]
+    pub fn with_exact_sizes() {
+        let mut buffer = [0; 30];
+        let localhost: embedded_nal::IpAddr = "127.0.0.1".parse().unwrap();
+        let builder: ConfigBuilder<IpBroker> = ConfigBuilder::new(localhost.into(), &mut buffer)
+            .session_state(BufferConfig::Exactly(8))
+            .rx_buffer(BufferConfig::Exactly(4))
+            .tx_buffer(BufferConfig::Exactly(3));
+
+        let config = builder.build();
+        assert!(config.tx_buffer.len() == 3);
+        assert!(config.rx_buffer.len() == 4);
+        assert!(config.state_buffer.len() == 8);
+    }
+
+    #[test]
+    pub fn with_max() {
+        let mut buffer = [0; 30];
+        let localhost: embedded_nal::IpAddr = "127.0.0.1".parse().unwrap();
+        let builder: ConfigBuilder<IpBroker> = ConfigBuilder::new(localhost.into(), &mut buffer)
+            .session_state(BufferConfig::Maximum(8));
+
+        let config = builder.build();
+        assert!(config.state_buffer.len() == 8);
+        assert!(config.tx_buffer.len() == 11);
+        assert!(config.rx_buffer.len() == 11);
+    }
+
+    #[test]
+    pub fn with_min() {
+        let mut buffer = [0; 30];
+        let localhost: embedded_nal::IpAddr = "127.0.0.1".parse().unwrap();
+        let builder: ConfigBuilder<IpBroker> = ConfigBuilder::new(localhost.into(), &mut buffer)
+            .session_state(BufferConfig::Minimum(20));
+
+        let config = builder.build();
+        assert!(config.state_buffer.len() == 20);
+        assert!(config.tx_buffer.len() == 5);
+        assert!(config.rx_buffer.len() == 5);
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 9b00ab1..ca55f48 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -23,22 +23,19 @@
 //! Below is a sample snippet showing how this library is used.
 //!
 //! ```no_run
-//! use minimq::{Config, Minimq, Publication};
+//! use minimq::{ConfigBuilder, Minimq, Publication};
 //!
 //! // Construct an MQTT client with a maximum packet size of 256 bytes
 //! // and a maximum of 16 messages that are allowed to be "in flight".
 //! // Messages are "in flight" if QoS::AtLeastOnce has not yet been acknowledged (PUBACK)
 //! // or QoS::ExactlyOnce has not been completed (PUBCOMP).
 //! // Connect to a broker at localhost - Use a client ID of "test".
-//! let mut rx_buffer = [0; 256];
-//! let mut tx_buffer = [0; 256];
-//! let mut session = [0; 256];
+//! let mut buffer = [0; 256];
 //! let localhost: embedded_nal::IpAddr = "127.0.0.1".parse().unwrap();
 //! let mut mqtt: Minimq<'_, _, _, minimq::broker::IpBroker> = Minimq::new(
 //!         std_embedded_nal::Stack::default(),
 //!         std_embedded_time::StandardClock::default(),
-//!         Config::new(localhost.into(), &mut rx_buffer, &mut tx_buffer)
-//!             .session_state(&mut session)
+//!         ConfigBuilder::new(localhost.into(), &mut buffer)
 //!             .client_id("test").unwrap(),
 //!         );
 //!
@@ -83,7 +80,7 @@ mod varint;
 mod will;
 
 pub use broker::Broker;
-pub use config::Config;
+pub use config::ConfigBuilder;
 pub use properties::Property;
 pub use publication::{DeferredPublication, Publication};
 pub use reason_codes::ReasonCode;
diff --git a/src/mqtt_client.rs b/src/mqtt_client.rs
index 77bcdc3..92b0969 100644
--- a/src/mqtt_client.rs
+++ b/src/mqtt_client.rs
@@ -6,7 +6,7 @@ use crate::{
     session_state::SessionState,
     types::{Properties, TopicFilter, Utf8String},
     will::SerializedWill,
-    Config, Error, MinimqError, Property, ProtocolError, QoS, {debug, error, info, warn},
+    ConfigBuilder, Error, MinimqError, Property, ProtocolError, QoS, {debug, error, info, warn},
 };
 
 use core::convert::{TryFrom, TryInto};
@@ -694,19 +694,15 @@ impl<'buf, TcpStack: TcpClientStack, Clock: embedded_time::Clock, Broker: crate:
     /// # Returns
     /// A `Minimq` object that can be used for publishing messages, subscribing to topics, and
     /// managing the MQTT state.
-    pub fn new(network_stack: TcpStack, clock: Clock, config: Config<'buf, Broker>) -> Self {
+    pub fn new(network_stack: TcpStack, clock: Clock, config: ConfigBuilder<'buf, Broker>) -> Self {
+        let config = config.build();
+
         let session_state = SessionState::new(
             config.client_id,
             config.state_buffer,
             config.tx_buffer.len(),
         );
 
-        let will = if let Some(crate::config::WillState::Serialized(will)) = config.will {
-            Some(will)
-        } else {
-            None
-        };
-
         Minimq {
             client: MqttClient {
                 sm: StateMachine::new(ClientContext::new(
@@ -716,7 +712,7 @@ impl<'buf, TcpStack: TcpClientStack, Clock: embedded_time::Clock, Broker: crate:
                 )),
                 downgrade_qos: config.downgrade_qos,
                 broker: config.broker,
-                will,
+                will: config.will,
                 network: InterfaceHolder::new(network_stack, config.tx_buffer),
                 max_packet_size: config.rx_buffer.len(),
             },
diff --git a/src/ser/mod.rs b/src/ser/mod.rs
index 82ee09b..e41fb76 100644
--- a/src/ser/mod.rs
+++ b/src/ser/mod.rs
@@ -81,6 +81,7 @@ impl core::fmt::Display for Error {
 pub struct MqttSerializer<'a> {
     buf: &'a mut [u8],
     index: usize,
+    with_header: bool,
 }
 
 impl<'a> MqttSerializer<'a> {
@@ -92,6 +93,15 @@ impl<'a> MqttSerializer<'a> {
         Self {
             buf,
             index: MAX_FIXED_HEADER_SIZE,
+            with_header: true,
+        }
+    }
+
+    pub fn new_without_header(buf: &'a mut [u8]) -> Self {
+        Self {
+            buf,
+            index: 0,
+            with_header: false,
         }
     }
 
@@ -100,7 +110,9 @@ impl<'a> MqttSerializer<'a> {
     /// # Note
     /// This does not append any MQTT headers.
     pub fn finish(self) -> &'a mut [u8] {
-        &mut self.buf[MAX_FIXED_HEADER_SIZE..self.index]
+        assert!(!self.with_header);
+
+        &mut self.buf[..self.index]
     }
 
     /// Encode an MQTT control packet into a buffer.
diff --git a/src/types.rs b/src/types.rs
index 98dacbe..b959250 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -25,6 +25,26 @@ pub enum Properties<'a> {
     },
 }
 
+impl<'a> Properties<'a> {
+    /// The length in bytes of the serialized properties.
+    pub fn size(&self) -> usize {
+        // Properties in MQTTv5 must be prefixed with a variable-length integer denoting the size
+        // of the all of the properties in bytes.
+        match self {
+            Properties::Slice(props) => props.iter().map(|prop| prop.size()).sum(),
+            Properties::CorrelatedSlice {
+                correlation,
+                properties,
+            } => properties
+                .iter()
+                .chain([*correlation].iter())
+                .map(|prop| prop.size())
+                .sum(),
+            Properties::DataBlock(block) => block.len(),
+        }
+    }
+}
+
 /// Used to progressively iterate across binary property blocks, deserializing them along the way.
 pub struct PropertiesIter<'a> {
     props: &'a [u8],
@@ -80,30 +100,22 @@ impl<'a> core::iter::IntoIterator for &'a Properties<'a> {
 impl<'a> serde::Serialize for Properties<'a> {
     fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
         let mut item = serializer.serialize_struct("Properties", 0)?;
+        item.serialize_field("_len", &Varint(self.size() as u32))?;
 
         // Properties in MQTTv5 must be prefixed with a variable-length integer denoting the size
         // of the all of the properties in bytes.
         match self {
             Properties::Slice(props) => {
-                let property_length: usize = props.iter().map(|prop| prop.size()).sum();
-                item.serialize_field("_len", &Varint(property_length as u32))?;
                 item.serialize_field("_props", props)?;
             }
             Properties::CorrelatedSlice {
                 correlation,
                 properties,
             } => {
-                let property_length: usize = properties
-                    .iter()
-                    .chain([*correlation].iter())
-                    .map(|prop| prop.size())
-                    .sum();
-                item.serialize_field("_len", &Varint(property_length as u32))?;
                 item.serialize_field("_correlation", &correlation)?;
                 item.serialize_field("_props", properties)?;
             }
             Properties::DataBlock(block) => {
-                item.serialize_field("_len", &Varint(block.len() as u32))?;
                 item.serialize_field("_data", block)?;
             }
         }
diff --git a/src/will.rs b/src/will.rs
index b5a4650..f5fcd20 100644
--- a/src/will.rs
+++ b/src/will.rs
@@ -1,6 +1,7 @@
 use crate::{
     properties::{Property, PropertyIdentifier},
     types::{BinaryData, Properties, Utf8String},
+    varint::Varint,
     ProtocolError, QoS, Retain,
 };
 
@@ -68,7 +69,7 @@ impl<'a> Will<'a> {
             data: BinaryData(self.data),
         };
 
-        let mut serializer = crate::ser::MqttSerializer::new(buf);
+        let mut serializer = crate::ser::MqttSerializer::new_without_header(buf);
         message.serialize(&mut serializer)?;
         Ok(SerializedWill {
             qos: self.qos,
@@ -77,6 +78,17 @@ impl<'a> Will<'a> {
         })
     }
 
+    /// Precalculate the length of the serialized will.
+    pub(crate) fn serialized_len(&self) -> usize {
+        let prop_len = {
+            let prop_size = Properties::Slice(self.properties).size();
+            Varint(prop_size as u32).len() + prop_size
+        };
+        let topic_len = self.topic.len() + core::mem::size_of::<u16>();
+        let payload_len = self.data.len() + core::mem::size_of::<u16>();
+        topic_len + payload_len + prop_len
+    }
+
     /// Set the retained status of the will.
     ///
     /// # Args


Test patch:
diff --git a/tests/at_least_once_subscription.rs b/tests/at_least_once_subscription.rs
index 3e74bdc..8c853ab 100644
--- a/tests/at_least_once_subscription.rs
+++ b/tests/at_least_once_subscription.rs
@@ -10,17 +10,13 @@ use std_embedded_time::StandardClock;
 fn main() -> std::io::Result<()> {
     env_logger::init();
 
-    let mut rx_buffer = [0u8; 256];
-    let mut tx_buffer = [0u8; 256];
-    let mut session = [0u8; 256];
-    let stack = std_embedded_nal::Stack::default();
+    let mut buffer = [0u8; 1024];
+    let stack = std_embedded_nal::Stack;
     let localhost = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
     let mut mqtt: Minimq<'_, _, _, minimq::broker::IpBroker> = Minimq::new(
         stack,
         StandardClock::default(),
-        minimq::Config::new(localhost.into(), &mut rx_buffer, &mut tx_buffer)
-            .session_state(&mut session)
-            .keepalive_interval(60),
+        minimq::ConfigBuilder::new(localhost.into(), &mut buffer).keepalive_interval(60),
     );
 
     let mut published = false;
diff --git a/tests/deferred_payload.rs b/tests/deferred_payload.rs
index 8d201d5..c9394d9 100644
--- a/tests/deferred_payload.rs
+++ b/tests/deferred_payload.rs
@@ -7,17 +7,13 @@ use std_embedded_time::StandardClock;
 fn main() -> std::io::Result<()> {
     env_logger::init();
 
-    let mut rx_buffer = [0u8; 256];
-    let mut tx_buffer = [0u8; 256];
-    let mut session = [0u8; 256];
+    let mut buffer = [0u8; 1024];
     let stack = std_embedded_nal::Stack;
     let localhost = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
     let mut mqtt: Minimq<'_, _, _, minimq::broker::IpBroker> = Minimq::new(
         stack,
         StandardClock::default(),
-        minimq::Config::new(localhost.into(), &mut rx_buffer, &mut tx_buffer)
-            .session_state(&mut session)
-            .keepalive_interval(60),
+        minimq::ConfigBuilder::new(localhost.into(), &mut buffer).keepalive_interval(60),
     );
 
     while !mqtt.client().is_connected() {
diff --git a/tests/exactly_once_subscription.rs b/tests/exactly_once_subscription.rs
index 60d3b80..5803566 100644
--- a/tests/exactly_once_subscription.rs
+++ b/tests/exactly_once_subscription.rs
@@ -10,17 +10,13 @@ use std_embedded_time::StandardClock;
 fn main() -> std::io::Result<()> {
     env_logger::init();
 
-    let mut rx_buffer = [0u8; 256];
-    let mut tx_buffer = [0u8; 256];
-    let mut session = [0u8; 256];
-    let stack = std_embedded_nal::Stack::default();
+    let mut buffer = [0u8; 1024];
+    let stack = std_embedded_nal::Stack;
     let localhost = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
     let mut mqtt: Minimq<'_, _, _, minimq::broker::IpBroker> = Minimq::new(
         stack,
         StandardClock::default(),
-        minimq::Config::new(localhost.into(), &mut rx_buffer, &mut tx_buffer)
-            .session_state(&mut session)
-            .keepalive_interval(60),
+        minimq::ConfigBuilder::new(localhost.into(), &mut buffer).keepalive_interval(60),
     );
 
     let mut published = false;
diff --git a/tests/exactly_once_transmission.rs b/tests/exactly_once_transmission.rs
index 4190635..f0c57c7 100644
--- a/tests/exactly_once_transmission.rs
+++ b/tests/exactly_once_transmission.rs
@@ -7,17 +7,13 @@ use std_embedded_time::StandardClock;
 fn main() -> std::io::Result<()> {
     env_logger::init();
 
-    let mut rx_buffer = [0u8; 256];
-    let mut tx_buffer = [0u8; 256];
-    let mut session = [0u8; 256];
-    let stack = std_embedded_nal::Stack::default();
+    let mut buffer = [0u8; 1024];
+    let stack = std_embedded_nal::Stack;
     let localhost = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
     let mut mqtt: Minimq<'_, _, _, minimq::broker::IpBroker> = Minimq::new(
         stack,
         StandardClock::default(),
-        minimq::Config::new(localhost.into(), &mut rx_buffer, &mut tx_buffer)
-            .session_state(&mut session)
-            .keepalive_interval(60),
+        minimq::ConfigBuilder::new(localhost.into(), &mut buffer).keepalive_interval(60),
     );
 
     let mut published = false;
diff --git a/tests/integration_test.rs b/tests/integration_test.rs
index 237520b..1de1e77 100644
--- a/tests/integration_test.rs
+++ b/tests/integration_test.rs
@@ -9,18 +9,14 @@ fn main() -> std::io::Result<()> {
 
     let will = Will::new("exit", "Test complete".as_bytes(), &[]).unwrap();
 
-    let mut rx_buffer = [0u8; 256];
-    let mut tx_buffer = [0u8; 256];
-    let mut session = [0u8; 256];
-    let mut will_buffer = [0; 64];
-    let stack = std_embedded_nal::Stack::default();
+    let mut buffer = [0u8; 1024];
+    let stack = std_embedded_nal::Stack;
     let localhost = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
     let mut mqtt: Minimq<'_, _, _, minimq::broker::IpBroker> = Minimq::new(
         stack,
         StandardClock::default(),
-        minimq::Config::new(localhost.into(), &mut rx_buffer, &mut tx_buffer)
-            .session_state(&mut session)
-            .will_with_buffer(&mut will_buffer, will)
+        minimq::ConfigBuilder::new(localhost.into(), &mut buffer)
+            .will(will)
             .unwrap()
             .keepalive_interval(60),
     );
diff --git a/tests/poll_result.rs b/tests/poll_result.rs
index 97d5beb..772284a 100644
--- a/tests/poll_result.rs
+++ b/tests/poll_result.rs
@@ -7,17 +7,13 @@ use std_embedded_time::StandardClock;
 fn main() -> std::io::Result<()> {
     env_logger::init();
 
-    let mut rx_buffer = [0u8; 256];
-    let mut tx_buffer = [0u8; 256];
-    let mut session = [0u8; 256];
-    let stack = std_embedded_nal::Stack::default();
+    let mut buffer = [0u8; 1024];
+    let stack = std_embedded_nal::Stack;
     let localhost = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
     let mut mqtt: Minimq<'_, _, _, minimq::broker::IpBroker> = Minimq::new(
         stack,
         StandardClock::default(),
-        minimq::Config::new(localhost.into(), &mut rx_buffer, &mut tx_buffer)
-            .session_state(&mut session)
-            .keepalive_interval(60),
+        minimq::ConfigBuilder::new(localhost.into(), &mut buffer).keepalive_interval(60),
     );
 
     let mut published = false;
diff --git a/tests/reconnection.rs b/tests/reconnection.rs
index 198c021..e281388 100644
--- a/tests/reconnection.rs
+++ b/tests/reconnection.rs
@@ -12,18 +12,15 @@ mod stack;
 fn main() -> std::io::Result<()> {
     env_logger::init();
 
-    let mut rx_buffer = [0u8; 256];
-    let mut tx_buffer = [0u8; 256];
-    let mut session = [0u8; 256];
     let sockets = std::cell::RefCell::new(Vec::new());
     let stack = stack::MitmStack::new(&sockets);
+
+    let mut buffer = [0u8; 1024];
     let localhost = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
     let mut mqtt: Minimq<'_, _, _, minimq::broker::IpBroker> = Minimq::new(
         stack,
         StandardClock::default(),
-        minimq::Config::new(localhost.into(), &mut rx_buffer, &mut tx_buffer)
-            .session_state(&mut session)
-            .keepalive_interval(1),
+        minimq::ConfigBuilder::new(localhost.into(), &mut buffer).keepalive_interval(1),
     );
 
     // 1. Poll until we're connected and subscribed to a test topic