add server pub/sub tests, fix pub/sub response (#27)

mini-redis server responses on PUB/SUB commands did not match real
redis.

src/client.rs

@@ -143,7 +143,7 @@ impl Client {
             let response = self.read_response().await?;
             match response {
                 Frame::Array(ref frame) => match frame.as_slice() {
-                    [subscribe, schannel]
+                    [subscribe, schannel, ..]
                         if subscribe.to_string() == "subscribe"
                             && &schannel.to_string() == channel =>
                     {
@@ -235,7 +235,7 @@ impl Subscriber {
             let response = self.read_response().await?;
             match response {
                 Frame::Array(ref frame) => match frame.as_slice() {
-                    [subscribe, schannel]
+                    [subscribe, schannel, ..]
                         if &subscribe.to_string() == "subscribe"
                             && &schannel.to_string() == channel =>
                     {
@@ -277,7 +277,7 @@ impl Subscriber {
             let response = self.read_response().await?;
             match response {
                 Frame::Array(ref frame) => match frame.as_slice() {
-                    [unsubscribe, uchannel] if &unsubscribe.to_string() == "unsubscribe" => {
+                    [unsubscribe, uchannel, ..] if &unsubscribe.to_string() == "unsubscribe" => {
                         //unsubscribed channel should exist in the subscribed list at this point
                         if self.subscribed_channels.remove(&uchannel.to_string()) == false {
                             return Err(response.to_error());

src/cmd/mod.rs

@@ -26,7 +26,7 @@ pub(crate) enum Command {
 }
 
 impl Command {
-    pub(crate) fn from_frame(frame: Frame) -> Result<Command, ParseError> {
+    pub(crate) fn from_frame(frame: Frame) -> crate::Result<Command> {
         let mut parse = Parse::new(frame)?;
 
         let command_name = parse.next_string()?.to_lowercase();

src/cmd/subscribe.rs

@@ -68,6 +68,7 @@ impl Subscribe {
                 let mut response = Frame::array();
                 response.push_bulk(Bytes::from_static(b"subscribe"));
                 response.push_bulk(Bytes::copy_from_slice(channel.as_bytes()));
+                response.push_int(subscriptions.len().saturating_add(1) as u64);
 
                 // Subscribe to channel
                 let rx = db.subscribe(channel.clone());
@@ -130,6 +131,7 @@ impl Subscribe {
                                 let mut response = Frame::array();
                                 response.push_bulk(Bytes::from_static(b"unsubscribe"));
                                 response.push_bulk(Bytes::copy_from_slice(channel.as_bytes()));
+                                response.push_int(subscriptions.len() as u64);
 
                                 dst.write_frame(&response).await?;
                             }

src/frame.rs

@@ -21,7 +21,7 @@ pub(crate) enum Error {
     Incomplete,
 
     /// Invalid message encoding
-    Invalid,
+    Other(crate::Error),
 }
 
 impl Frame {
@@ -30,7 +30,7 @@ impl Frame {
         Frame::Array(vec![])
     }
 
-    /// Push a "bulk" frame into the array. `self` must be an Array frame
+    /// Push a "bulk" frame into the array. `self` must be an Array frame.
     ///
     /// # Panics
     ///
@@ -44,6 +44,20 @@ impl Frame {
         }
     }
 
+    /// Push an "integer" frame into the array. `self` must be an Array frame.
+    ///
+    /// # Panics
+    ///
+    /// panics if `self` is not an array
+    pub(crate) fn push_int(&mut self, value: u64) {
+        match self {
+            Frame::Array(vec) => {
+                vec.push(Box::new(Frame::Integer(value)));
+            }
+            _ => panic!("not an array frame"),
+        }
+    }
+
     /// Checks if an entire message can be decoded from `src`
     pub(crate) fn check(src: &mut Cursor<&[u8]>) -> Result<(), Error> {
         match get_u8(src)? {
@@ -80,7 +94,7 @@ impl Frame {
 
                 Ok(())
             }
-            _ => Err(Error::Invalid),
+            actual => Err(format!("protocol error; invalid frame type byte `{}`", actual).into()),
         }
     }
 
@@ -114,7 +128,7 @@ impl Frame {
                     let line = get_line(src)?;
 
                     if line != b"-1" {
-                        return Err(Error::Invalid);
+                        return Err("protocol error; invalid frame format".into());
                     }
 
                     Ok(Frame::Null)
@@ -213,7 +227,7 @@ fn get_decimal(src: &mut Cursor<&[u8]>) -> Result<u64, Error> {
 
     let line = get_line(src)?;
 
-    atoi::<u64>(line).ok_or(Error::Invalid)
+    atoi::<u64>(line).ok_or_else(|| "protocol error; invalid frame format".into())
 }
 
 /// Find a line
@@ -236,15 +250,27 @@ fn get_line<'a>(src: &mut Cursor<&'a [u8]>) -> Result<&'a [u8], Error> {
     Err(Error::Incomplete)
 }
 
+impl From<String> for Error {
+    fn from(src: String) -> Error {
+        Error::Other(src.into())
+    }
+}
+
+impl From<&str> for Error {
+    fn from(src: &str) -> Error {
+        src.to_string().into()
+    }
+}
+
 impl From<FromUtf8Error> for Error {
     fn from(_src: FromUtf8Error) -> Error {
-        unimplemented!();
+        "protocol error; invalid frame format".into()
     }
 }
 
 impl From<TryFromIntError> for Error {
     fn from(_src: TryFromIntError) -> Error {
-        unimplemented!();
+        "protocol error; invalid frame format".into()
     }
 }
 
@@ -254,7 +280,7 @@ impl fmt::Display for Error {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Error::Incomplete => "stream ended early".fmt(fmt),
-            Error::Invalid => "invalid frame format".fmt(fmt),
+            Error::Other(err) => err.fmt(fmt),
         }
     }
 }

src/parse.rs

@@ -1,7 +1,7 @@
 use crate::Frame;
 
 use bytes::Bytes;
-use std::{error, fmt, io, str, vec};
+use std::{fmt, str, vec};
 
 /// Utility for parsing a command
 #[derive(Debug)]
@@ -12,14 +12,14 @@ pub(crate) struct Parse {
 #[derive(Debug)]
 pub(crate) enum ParseError {
     EndOfStream,
-    Invalid,
+    Other(crate::Error),
 }
 
 impl Parse {
     pub(crate) fn new(frame: Frame) -> Result<Parse, ParseError> {
         let array = match frame {
             Frame::Array(array) => array,
-            _ => return Err(ParseError::Invalid),
+            frame => return Err(format!("protocol error; expected array, got {:?}", frame).into()),
         };
 
         Ok(Parse {
@@ -39,8 +39,8 @@ impl Parse {
             Frame::Simple(s) => Ok(s),
             Frame::Bulk(data) => str::from_utf8(&data[..])
                 .map(|s| s.to_string())
-                .map_err(|_| ParseError::Invalid),
+                .map_err(|_| "protocol error; invalid string".into()),
-            _ => Err(ParseError::Invalid),
+            frame => Err(format!("protocol error; expected simple frame or bulk frame, got {:?}", frame).into()),
         }
     }
 
@@ -48,18 +48,20 @@ impl Parse {
         match self.next()? {
             Frame::Simple(s) => Ok(Bytes::from(s.into_bytes())),
             Frame::Bulk(data) => Ok(data),
-            _ => Err(ParseError::Invalid),
+            frame => Err(format!("protocol error; expected simple frame or bulk frame, got {:?}", frame).into()),
         }
     }
 
     pub(crate) fn next_int(&mut self) -> Result<u64, ParseError> {
         use atoi::atoi;
 
+        const MSG: &str = "protocol error; invalid number";
+
         match self.next()? {
             Frame::Integer(v) => Ok(v),
-            Frame::Simple(data) => atoi::<u64>(data.as_bytes()).ok_or(ParseError::Invalid),
+            Frame::Simple(data) => atoi::<u64>(data.as_bytes()).ok_or_else(|| MSG.into()),
-            Frame::Bulk(data) => atoi::<u64>(&data).ok_or(ParseError::Invalid),
+            Frame::Bulk(data) => atoi::<u64>(&data).ok_or_else(|| MSG.into()),
-            _ => Err(ParseError::Invalid),
+            frame => Err(format!("protocol error; expected int frame but got {:?}", frame).into()),
         }
     }
 
@@ -68,29 +70,33 @@ impl Parse {
         if self.parts.next().is_none() {
             Ok(())
         } else {
-            Err(ParseError::Invalid)
+            Err("protocol error; expected end of frame, but there was more".into())
         }
     }
 }
 
-impl From<ParseError> for io::Error {
+impl From<String> for ParseError {
-    fn from(src: ParseError) -> io::Error {
+    fn from(src: String) -> ParseError {
-        io::Error::new(io::ErrorKind::Other, format!("{}", src))
+        ParseError::Other(src.into())
+    }
+}
+
+impl From<&str> for ParseError {
+    fn from(src: &str) -> ParseError {
+        src.to_string().into()
     }
 }
 
 impl fmt::Display for ParseError {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let msg = match self {
+        match self {
-            ParseError::EndOfStream => "end of stream".to_string(),
+            ParseError::EndOfStream => {
-            ParseError::Invalid => "invalid".to_string(),
+                "protocol error; unexpected end of stream".fmt(f)
-        };
+            }
-        write!(f, "{}", &msg)
+            ParseError::Other(err) => err.fmt(f),
+        }
     }
 }
 
 impl std::error::Error for ParseError {
-    fn source(&self) -> Option<&(dyn error::Error + 'static)> {
-        None
-    }
 }

tests/server.rs

@@ -3,7 +3,6 @@ use mini_redis::server;
 use std::net::SocketAddr;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::{TcpListener, TcpStream};
-use tokio::task::JoinHandle;
 use tokio::time::{self, Duration};
 
 /// A basic "hello world" style test. A server instance is started in a
@@ -12,7 +11,7 @@ use tokio::time::{self, Duration};
 /// level.
 #[tokio::test]
 async fn key_value_get_set() {
-    let (addr, _handle) = start_server().await;
+    let addr = start_server().await;
 
     // Establish a connection to the server
     let mut stream = TcpStream::connect(addr).await.unwrap();
@@ -22,17 +21,15 @@ async fn key_value_get_set() {
 
     // Read nil response
     let mut response = [0; 5];
-
     stream.read_exact(&mut response).await.unwrap();
-
     assert_eq!(b"$-1\r\n", &response);
 
     // Set a key
     stream.write_all(b"*3\r\n$3\r\nSET\r\n$5\r\nhello\r\n$5\r\nworld\r\n").await.unwrap();
 
     // Read OK
+    let mut response = [0; 5];
     stream.read_exact(&mut response).await.unwrap();
-
     assert_eq!(b"+OK\r\n", &response);
 
     // Get the key, data is present
@@ -40,9 +37,7 @@ async fn key_value_get_set() {
 
     // Read "world" response
     let mut response = [0; 11];
-
     stream.read_exact(&mut response).await.unwrap();
-
     assert_eq!(b"$5\r\nworld\r\n", &response);
 }
 
@@ -58,7 +53,7 @@ async fn key_value_get_set() {
 async fn key_value_timeout() {
     tokio::time::pause();
 
-    let (addr, _handle) = start_server().await;
+    let addr = start_server().await;
 
     // Establish a connection to the server
     let mut stream = TcpStream::connect(addr).await.unwrap();
@@ -98,40 +93,153 @@ async fn key_value_timeout() {
     assert_eq!(b"$-1\r\n", &response);
 }
 
-// In this case we test that server responds acurately to
-// SUBSCRIBE and UNSUBSCRIBE commands
 #[tokio::test]
-async fn subscribe_unsubscribe() {
+async fn pub_sub() {
-    let (addr, _handle) = start_server().await;
+    let addr = start_server().await;
 
-    let mut stream = TcpStream::connect(addr).await.unwrap();
+    let mut publisher = TcpStream::connect(addr).await.unwrap();
 
-    // send SUBSCRIBE command
+    // Publish a message, there are no subscribers yet so the server will
-    stream.write_all(b"*2\r\n$9\r\nsubscribe\r\n$5\r\nhello\r\n").await.unwrap();
+    // return `0`.
+    publisher.write_all(b"*3\r\n$7\r\nPUBLISH\r\n$5\r\nhello\r\n$5\r\nworld\r\n").await.unwrap();
 
-    // Read response
+    let mut response = [0; 4];
-    let mut response = [0; 30];
+    publisher.read_exact(&mut response).await.unwrap();
+    assert_eq!(b":0\r\n", &response);
 
-    stream.read_exact(&mut response).await.unwrap();
+    // Create a subscriber. This subscriber will only subscribe to the `hello`
+    // channel.
+    let mut sub1 = TcpStream::connect(addr).await.unwrap();
+    sub1.write_all(b"*2\r\n$9\r\nSUBSCRIBE\r\n$5\r\nhello\r\n").await.unwrap();
 
-    assert_eq!(b"*2\r\n$9\r\nsubscribe\r\n$5\r\nhello\r\n", &response);
+    // Read the subscribe response
+    let mut response = [0; 34];
+    sub1.read_exact(&mut response).await.unwrap();
+    assert_eq!(&b"*3\r\n$9\r\nsubscribe\r\n$5\r\nhello\r\n:1\r\n"[..], &response[..]);
 
-    // send UNSUBSCRIBE command
+    // Publish a message, there now is a subscriber
-    stream.write_all(b"*2\r\n$11\r\nunsubscribe\r\n$5\r\nhello\r\n").await.unwrap();
+    publisher.write_all(b"*3\r\n$7\r\nPUBLISH\r\n$5\r\nhello\r\n$5\r\nworld\r\n").await.unwrap();
 
-    let mut response = [0; 33];
+    let mut response = [0; 4];
+    publisher.read_exact(&mut response).await.unwrap();
+    assert_eq!(b":1\r\n", &response);
 
-    stream.read_exact(&mut response).await.unwrap();
+    // The first subscriber received the message
+    let mut response = [0; 39];
+    sub1.read_exact(&mut response).await.unwrap();
+    assert_eq!(&b"*3\r\n$7\r\nmessage\r\n$5\r\nhello\r\n$5\r\nworld\r\n"[..], &response[..]);
 
-    assert_eq!(b"*2\r\n$11\r\nunsubscribe\r\n", &response[0..22]);
+    // Create a second subscriber
-    assert_eq!(b"$5\r\nhello\r\n", &response[22..33]);
+    //
+    // This subscriber will be subscribed to both `hello` and `foo`
+    let mut sub2 = TcpStream::connect(addr).await.unwrap();
+    sub2.write_all(b"*3\r\n$9\r\nSUBSCRIBE\r\n$5\r\nhello\r\n$3\r\nfoo\r\n").await.unwrap();
+
+    // Read the subscribe response
+    let mut response = [0; 34];
+    sub2.read_exact(&mut response).await.unwrap();
+    assert_eq!(&b"*3\r\n$9\r\nsubscribe\r\n$5\r\nhello\r\n:1\r\n"[..], &response[..]);
+    let mut response = [0; 32];
+    sub2.read_exact(&mut response).await.unwrap();
+    assert_eq!(&b"*3\r\n$9\r\nsubscribe\r\n$3\r\nfoo\r\n:2\r\n"[..], &response[..]);
+
+    // Publish another message on `hello`, there are two subscribers
+    publisher.write_all(b"*3\r\n$7\r\nPUBLISH\r\n$5\r\nhello\r\n$5\r\njazzy\r\n").await.unwrap();
+
+    let mut response = [0; 4];
+    publisher.read_exact(&mut response).await.unwrap();
+    assert_eq!(b":2\r\n", &response);
+
+    // Publish a message on `foo`, there is only one subscriber
+    publisher.write_all(b"*3\r\n$7\r\nPUBLISH\r\n$3\r\nfoo\r\n$3\r\nbar\r\n").await.unwrap();
+
+    let mut response = [0; 4];
+    publisher.read_exact(&mut response).await.unwrap();
+    assert_eq!(b":1\r\n", &response);
+
+    // The first subscriber received the message
+    let mut response = [0; 39];
+    sub1.read_exact(&mut response).await.unwrap();
+    assert_eq!(&b"*3\r\n$7\r\nmessage\r\n$5\r\nhello\r\n$5\r\njazzy\r\n"[..], &response[..]);
+
+    // The second subscriber received the message
+    let mut response = [0; 39];
+    sub2.read_exact(&mut response).await.unwrap();
+    assert_eq!(&b"*3\r\n$7\r\nmessage\r\n$5\r\nhello\r\n$5\r\njazzy\r\n"[..], &response[..]);
+
+    // The first subscriber **did not** receive the second message
+    let mut response = [0; 1];
+    time::timeout(Duration::from_millis(100), sub1.read(&mut response)).await.unwrap_err();
+
+    // The second subscriber **did** receive the message
+    let mut response = [0; 35];
+    sub2.read_exact(&mut response).await.unwrap();
+    assert_eq!(&b"*3\r\n$7\r\nmessage\r\n$3\r\nfoo\r\n$3\r\nbar\r\n"[..], &response[..]);
+}
+
+#[tokio::test]
+async fn manage_subscription() {
+    let addr = start_server().await;
+
+    let mut publisher = TcpStream::connect(addr).await.unwrap();
+
+    // Create a subscriber
+    let mut sub = TcpStream::connect(addr).await.unwrap();
+    sub.write_all(b"*2\r\n$9\r\nSUBSCRIBE\r\n$5\r\nhello\r\n").await.unwrap();
+
+    // Read the subscribe response
+    let mut response = [0; 34];
+    sub.read_exact(&mut response).await.unwrap();
+    assert_eq!(&b"*3\r\n$9\r\nsubscribe\r\n$5\r\nhello\r\n:1\r\n"[..], &response[..]);
+
+    // Update subscription to add `foo`
+    sub.write_all(b"*2\r\n$9\r\nSUBSCRIBE\r\n$3\r\nfoo\r\n").await.unwrap();
+
+    let mut response = [0; 32];
+    sub.read_exact(&mut response).await.unwrap();
+    assert_eq!(&b"*3\r\n$9\r\nsubscribe\r\n$3\r\nfoo\r\n:2\r\n"[..], &response[..]);
+
+    // Update subscription to remove `hello`
+    sub.write_all(b"*2\r\n$11\r\nUNSUBSCRIBE\r\n$5\r\nhello\r\n").await.unwrap();
+
+    let mut response = [0; 37];
+    sub.read_exact(&mut response).await.unwrap();
+    assert_eq!(&b"*3\r\n$11\r\nunsubscribe\r\n$5\r\nhello\r\n:1\r\n"[..], &response[..]);
+
+    // Publish a message to `hello` and then a message to `foo`
+    publisher.write_all(b"*3\r\n$7\r\nPUBLISH\r\n$5\r\nhello\r\n$5\r\nworld\r\n").await.unwrap();
+    let mut response = [0; 4];
+    publisher.read_exact(&mut response).await.unwrap();
+    assert_eq!(b":0\r\n", &response);
+
+    publisher.write_all(b"*3\r\n$7\r\nPUBLISH\r\n$3\r\nfoo\r\n$3\r\nbar\r\n").await.unwrap();
+    let mut response = [0; 4];
+    publisher.read_exact(&mut response).await.unwrap();
+    assert_eq!(b":1\r\n", &response);
+
+    // Receive the message
+    // The second subscriber **did** receive the message
+    let mut response = [0; 35];
+    sub.read_exact(&mut response).await.unwrap();
+    assert_eq!(&b"*3\r\n$7\r\nmessage\r\n$3\r\nfoo\r\n$3\r\nbar\r\n"[..], &response[..]);
+
+    // No more messages
+    let mut response = [0; 1];
+    time::timeout(Duration::from_millis(100), sub.read(&mut response)).await.unwrap_err();
+
+    // Unsubscribe from all channels
+    sub.write_all(b"*1\r\n$11\r\nunsubscribe\r\n").await.unwrap();
+
+    let mut response = [0; 35];
+    sub.read_exact(&mut response).await.unwrap();
+    assert_eq!(&b"*3\r\n$11\r\nunsubscribe\r\n$3\r\nfoo\r\n:0\r\n"[..], &response[..]);
 }
 
 // In this case we test that server Responds with an Error message if a client
 // sends an unknown command
 #[tokio::test]
 async fn send_error_unknown_command() {
-    let (addr, _handle) = start_server().await;
+    let addr = start_server().await;
 
     // Establish a connection to the server
     let mut stream = TcpStream::connect(addr).await.unwrap();
@@ -150,7 +258,7 @@ async fn send_error_unknown_command() {
 // sends an GET or SET command after a SUBSCRIBE
 #[tokio::test]
 async fn send_error_get_set_after_subscribe() {
-    let (addr, _handle) = start_server().await;
+    let addr = start_server().await;
 
     let mut stream = TcpStream::connect(addr).await.unwrap();
 
@@ -178,13 +286,13 @@ async fn send_error_get_set_after_subscribe() {
     assert_eq!(b"-ERR unknown command \'get\'\r\n", &response);
 }
 
-async fn start_server() -> (SocketAddr, JoinHandle<mini_redis::Result<()>>) {
+async fn start_server() -> SocketAddr {
     let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
     let addr = listener.local_addr().unwrap();
 
-    let handle = tokio::spawn(async move {
+    tokio::spawn(async move {
         server::run(listener, tokio::signal::ctrl_c()).await
     });
 
-    (addr, handle)
+    addr
 }