lostecho/mini-redis
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
mini-redis server responses on PUB/SUB commands did not match real
redis.
This commit is contained in:
Carl Lerche
2020-04-06 16:34:12 -07:00
committed by GitHub
parent 922919a9d4
commit d4f0dac671
6 changed files with 205 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/client.rs
View File

@@ -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());

2
src/cmd/mod.rs
View File

@@ -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();

2
src/cmd/subscribe.rs
View File

@@ -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?;
}

42
src/frame.rs
View File

@@ -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),
}
}
}

48
src/parse.rs
View File

@@ -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),
_ => Err(ParseError::Invalid),
.map_err(|_| "protocol error; invalid string".into()),
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::Bulk(data) => atoi::<u64>(&data).ok_or(ParseError::Invalid),
_ => Err(ParseError::Invalid),
Frame::Simple(data) => atoi::<u64>(data.as_bytes()).ok_or_else(|| MSG.into()),
Frame::Bulk(data) => atoi::<u64>(&data).ok_or_else(|| MSG.into()),
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 {
fn from(src: ParseError) -> io::Error {
io::Error::new(io::ErrorKind::Other, format!("{}", src))
impl From<String> for ParseError {
fn from(src: String) -> ParseError {
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 {
ParseError::EndOfStream => "end of stream".to_string(),
ParseError::Invalid => "invalid".to_string(),
};
write!(f, "{}", &msg)
match self {
ParseError::EndOfStream => {
"protocol error; unexpected end of stream".fmt(f)
}
ParseError::Other(err) => err.fmt(f),
}
}
}
impl std::error::Error for ParseError {
fn source(&self) -> Option<&(dyn error::Error + 'static)> {
None
}
}

168
tests/server.rs
View File

@@ -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() {
let (addr, _handle) = start_server().await;
async fn pub_sub() {
let addr = start_server().await;
let mut stream = TcpStream::connect(addr).await.unwrap();
let mut publisher = TcpStream::connect(addr).await.unwrap();
// send SUBSCRIBE command
stream.write_all(b"*2\r\n$9\r\nsubscribe\r\n$5\r\nhello\r\n").await.unwrap();
// Publish a message, there are no subscribers yet so the server will
// 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; 30];
let mut response = [0; 4];
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
stream.write_all(b"*2\r\n$11\r\nunsubscribe\r\n$5\r\nhello\r\n").await.unwrap();
// Publish a message, there now is a subscriber
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]);
assert_eq!(b"$5\r\nhello\r\n", &response[22..33]);
// Create a second subscriber
//
// 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
}