<commit_before>/*
 * Copyright 2017 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
 * http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
 * <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
 * option. This file may not be copied, modified, or distributed
 * except according to those terms.
 */

//! An implementation of the RESP protocol

use std::io;
use std::str;

use bytes::{BufMut, BytesMut};

use tokio_io::codec::{Decoder, Encoder};

use super::error;
use super::error::Error;

/// A single RESP value, this owns the data that is read/to-be written to Redis.
///
/// It is cloneable to allow multiple copies to be delivered in certain circumstances, e.g. multiple
/// subscribers to the same topic. TODO - revisit this assumption to make sure its sound, perhaps a
/// single-subscriber enforcement would make more sense, or sharing via `Arc`
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum RespValue {
    /// Zero, one or more other `RespValue`s.
    Array(Vec<RespValue>),

    /// A bulk string.  In Redis terminology a string is a byte-array, so this is stored as a
    /// vector of `u8`s to allow clients to interpret the bytes as appropriate.
    BulkString(Vec<u8>),

    /// An error from the Redis server
    Error(String),

    Integer(usize),

    SimpleString(String),
}

/// A trait to be implemented for every time which can be read from a RESP value.
///
/// Implementing this trait on a type means that type becomes a valid return type for calls such as `send` on
/// `client::PairedConnection`
pub trait FromResp: Sized {
    /// Return a `Result` containing either `Self` or `Error`.  Errors can occur due to either: a) the particular
    /// `RespValue` being incompatible with the required type, or b) a remote Redis error occuring.
    fn from_resp(resp: RespValue) -> Result<Self, Error>;
}

impl FromResp for RespValue {
    fn from_resp(resp: RespValue) -> Result<RespValue, Error> {
        match resp {
            // TODO - this logic is repeated in all implementations, move somewhere central
            RespValue::Error(string) => Err(Error::Remote(string)),
            x => Ok(x),
        }
    }
}

impl FromResp for String {
    fn from_resp(resp: RespValue) -> Result<String, Error> {
        match resp {
            // TODO - the "cannot be converted" probably needs to explain what cannot be converted
            RespValue::Array(_) => Err(error::resp("Cannot convert into a string", resp)),
            RespValue::BulkString(ref bytes) => Ok(String::from_utf8_lossy(bytes).into_owned()),
            RespValue::Error(string) => Err(Error::Remote(string)),
            RespValue::Integer(i) => Ok(i.to_string()),
            RespValue::SimpleString(string) => Ok(string),
        }
    }
}

impl FromResp for usize {
    fn from_resp(resp: RespValue) -> Result<usize, Error> {
        match resp {
            RespValue::Error(string) => Err(Error::Remote(string)),
            RespValue::Integer(i) => Ok(i),
            _ => Err(error::resp("Cannot be converted into a usize", resp)),
        }
    }
}

impl FromResp for () {
    fn from_resp(resp: RespValue) -> Result<(), Error> {
        match resp {
            RespValue::Error(string) => Err(Error::Remote(string)),
            RespValue::SimpleString(string) => {
                match string.as_ref() {
                    "OK" => Ok(()),
                    _ => {
                        Err(Error::RESP(format!("Unexpected value within SimpleString: {}",
                                                string),
                                        None))
                    }
                }
            }
            _ => Err(error::resp("Unexpected value", resp)),
        }
    }
}

/// A trait to be implemented on types that can be automatically converted into `RespValue`s.
///
/// A `From<T>` where `T: ToResp` has been implemented so that everything that implements `ToResp`
/// can be converted with conversion traits.
pub trait ToResp {
    fn to_resp(&self) -> RespValue;
}

impl<'a, T> ToResp for &'a [T]
    where T: ToResp
{
    fn to_resp(&self) -> RespValue {
        RespValue::Array(self.as_ref().iter().map(|x| x.to_resp()).collect())
    }
}

impl<'a, T> ToResp for Vec<T>
    where T: ToResp
{
    fn to_resp(&self) -> RespValue {
        RespValue::Array(self.into_iter().map(|v| v.to_resp()).collect())
    }
}

impl<T: ToResp> From<T> for RespValue {
    fn from(from: T) -> RespValue {
        from.to_resp()
    }
}

/// A specific trait to convert into a `RespValue::BulkString`
pub trait ToRespString {
    fn to_resp_string(&self) -> RespValue;
}

impl ToRespString for String {
    fn to_resp_string(&self) -> RespValue {
        RespValue::BulkString(self.as_bytes().into())
    }
}

impl<'a> ToRespString for &'a str {
    fn to_resp_string(&self) -> RespValue {
        RespValue::BulkString(self.as_bytes().into())
    }
}

impl<T> ToResp for T
    where T: ToRespString
{
    fn to_resp(&self) -> RespValue {
        self.to_resp_string()
    }
}

/// Codec to read frames
pub struct RespCodec;

fn write_rn(buf: &mut BytesMut) {
    buf.put_u8(b'\r');
    buf.put_u8(b'\n');
}

fn check_and_reserve(buf: &mut BytesMut, amt: usize) {
    let remaining_bytes = buf.remaining_mut();
    if remaining_bytes < amt {
        buf.reserve(amt);
    }
}

fn write_header(symb: u8, len: usize, buf: &mut BytesMut) {
    let len_as_string = len.to_string();
    let len_as_bytes = len_as_string.as_bytes();
    let header_bytes = 1 + len_as_bytes.len() + 2;
    check_and_reserve(buf, header_bytes);
    buf.put_u8(symb);
    buf.extend(len_as_bytes);
    write_rn(buf);
}

fn write_simple_string(symb: u8, string: &str, buf: &mut BytesMut) {
    let bytes = string.as_bytes();
    let size = 1 + bytes.len() + 2;
    check_and_reserve(buf, size);
    buf.put_u8(symb);
    buf.extend(bytes);
    write_rn(buf);
}

impl Encoder for RespCodec {
    type Item = RespValue;
    type Error = io::Error;

    fn encode(&mut self, msg: RespValue, buf: &mut BytesMut) -> Result<(), Self::Error> {
        match msg {
            RespValue::Array(ary) => {
                write_header(b'*', ary.len(), buf);
                for v in ary {
                    self.encode(v, buf)?;
                }
            }
            RespValue::BulkString(bstr) => {
                let len = bstr.len();
                write_header(b'$', len, buf);
                check_and_reserve(buf, len + 2);
                buf.extend(bstr);
                write_rn(buf);
            }
            RespValue::Error(ref string) => {
                write_simple_string(b'-', string, buf);
            }
            RespValue::Integer(val) => {
                // Simple integer are just the header
                write_header(b':', val, buf);
            }
            RespValue::SimpleString(ref string) => {
                write_simple_string(b'+', string, buf);
            }
        }
        Ok(())
    }
}

#[inline]
fn parse_error(message: String) -> Error {
    Error::RESP(message, None)
}

/// Many RESP types have their length (which is either bytes or "number of elements", depending on context)
/// encoded as a string, terminated by "\r\n", this looks for them.
///
/// Only return the string if the whole sequence is complete, including the terminator bytes (but those final
/// two bytes will not be returned)
///
/// TODO - rename this function potentially, it's used for simple integers too
fn scan_integer<'a>(buf: &'a mut BytesMut, idx: usize) -> Result<Option<(usize, &'a [u8])>, Error> {
    let length = buf.len();
    let mut at_end = false;
    let mut pos = idx;
    loop {
        if length <= pos {
            return Ok(None);
        }
        match (at_end, buf[pos]) {
            (true, b'\n') => return Ok(Some((pos + 1, &buf[idx..pos - 1]))),
            (false, b'\r') => at_end = true,
            (false, b'0'...b'9') => (),
            (_, val) => return Err(parse_error(format!("Unexpected byte in size_string: {}", val))),
        }
        pos += 1;
    }
}

fn scan_string(buf: &mut BytesMut, idx: usize) -> Option<(usize, String)> {
    let length = buf.len();
    let mut at_end = false;
    let mut pos = idx;
    loop {
        if length <= pos {
            return None;
        }
        match (at_end, buf[pos]) {
            (true, b'\n') => {
                let value = String::from_utf8_lossy(&buf[idx..pos - 1]).into_owned();
                return Some((pos + 1, value));
            }
            (true, _) => at_end = false,
            (false, b'\r') => at_end = true,
            (false, _) => (),
        }
        pos += 1;
    }
}

fn decode_raw_integer(buf: &mut BytesMut, idx: usize) -> Result<Option<(usize, usize)>, Error> {
    match scan_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, int_str))) => {
            let int: usize = str::from_utf8(int_str)
                .expect("Not a string")
                .parse()
                .expect("Not an integer");
            Ok(Some((pos, int)))
        }
        Err(e) => Err(e),
    }
}

type DecodeResult = Result<Option<(usize, RespValue)>, Error>;

fn decode_bulk_string(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match decode_raw_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, size))) => {
            let remaining = buf.len() - pos;
            let required_bytes = size + 2;

            if remaining < required_bytes {
                return Ok(None);
            }

            let bulk_string = RespValue::BulkString(buf[pos..(pos + size)].to_vec());
            Ok(Some((pos + required_bytes, bulk_string)))
        }
        Err(e) => Err(e),
    }
}

fn decode_array(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match decode_raw_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, size))) => {
            let mut pos = pos;
            let mut values = Vec::with_capacity(size);
            for _ in 0..size {
                match decode(buf, pos) {
                    Ok(None) => return Ok(None),
                    Ok(Some((new_pos, value))) => {
                        values.push(value);
                        pos = new_pos;
                    }
                    Err(e) => return Err(e),
                }
            }
            Ok(Some((pos, RespValue::Array(values))))
        }
        Err(e) => Err(e),
    }
}

fn decode_integer(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match decode_raw_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, int))) => Ok(Some((pos, RespValue::Integer(int)))),
        Err(e) => Err(e),
    }
}

/// A simple string is any series of bytes that ends with `\r\n`
fn decode_simple_string(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match scan_string(buf, idx) {
        None => Ok(None),
        Some((pos, string)) => Ok(Some((pos, RespValue::SimpleString(string)))),
    }
}

fn decode_error(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match scan_string(buf, idx) {
        None => Ok(None),
        Some((pos, string)) => Ok(Some((pos, RespValue::Error(string)))),
    }
}

fn decode(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    let length = buf.len();
    if length <= idx {
        return Ok(None);
    }

    let first_byte = buf[idx];
    match first_byte {
        b'$' => decode_bulk_string(buf, idx + 1),
        b'*' => decode_array(buf, idx + 1),
        b':' => decode_integer(buf, idx + 1),
        b'+' => decode_simple_string(buf, idx + 1),
        b'-' => decode_error(buf, idx + 1),
        _ => Err(parse_error(format!("Unexpected byte: {}", first_byte))),
    }
}

impl Decoder for RespCodec {
    type Item = RespValue;
    type Error = Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        match decode(buf, 0) {
            Ok(None) => Ok(None),
            Ok(Some((pos, item))) => {
                buf.split_to(pos);
                Ok(Some(item))
            }
            Err(e) => Err(e),
        }
    }
}

#[cfg(test)]
mod tests {
    use bytes::BytesMut;

    use tokio_io::codec::{Decoder, Encoder};

    use super::{RespCodec, RespValue};

    #[test]
    fn test_bulk_string() {
        let resp_object = RespValue::BulkString("THISISATEST".as_bytes().to_vec());
        let mut bytes = BytesMut::new();
        let mut codec = RespCodec;
        codec.encode(resp_object.clone(), &mut bytes).unwrap();
        assert_eq!(b"$11\r\nTHISISATEST\r\n".to_vec(), bytes.to_vec());

        let deserialized = codec.decode(&mut bytes).unwrap().unwrap();
        assert_eq!(deserialized, resp_object);
    }

    #[test]
    fn test_array() {
        let resp_object = RespValue::Array(vec!["TEST1".into(), "TEST2".into()]);
        let mut bytes = BytesMut::new();
        let mut codec = RespCodec;
        codec.encode(resp_object.clone(), &mut bytes).unwrap();
        assert_eq!(b"*2\r\n$5\r\nTEST1\r\n$5\r\nTEST2\r\n".to_vec(),
                   bytes.to_vec());

        let deserialized = codec.decode(&mut bytes).unwrap().unwrap();
        assert_eq!(deserialized, resp_object);
    }
}<commit_msg>Parsing improvements<commit_after>/*
 * Copyright 2017 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
 * http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
 * <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
 * option. This file may not be copied, modified, or distributed
 * except according to those terms.
 */

//! An implementation of the RESP protocol

use std::io;
use std::str;

use bytes::{BufMut, BytesMut};

use tokio_io::codec::{Decoder, Encoder};

use super::error;
use super::error::Error;

/// A single RESP value, this owns the data that is read/to-be written to Redis.
///
/// It is cloneable to allow multiple copies to be delivered in certain circumstances, e.g. multiple
/// subscribers to the same topic. TODO - revisit this assumption to make sure its sound, perhaps a
/// single-subscriber enforcement would make more sense, or sharing via `Arc`
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum RespValue {
    /// Zero, one or more other `RespValue`s.
    Array(Vec<RespValue>),

    /// A bulk string.  In Redis terminology a string is a byte-array, so this is stored as a
    /// vector of `u8`s to allow clients to interpret the bytes as appropriate.
    BulkString(Vec<u8>),

    /// An error from the Redis server
    Error(String),

    Integer(usize),

    SimpleString(String),
}

impl RespValue {
    fn to_result(self) -> Result<RespValue, Error> {
        match self {
            RespValue::Error(string) => Err(Error::Remote(string)),
            x => Ok(x),
        }
    }
}

/// A trait to be implemented for every time which can be read from a RESP value.
///
/// Implementing this trait on a type means that type becomes a valid return type for calls such as `send` on
/// `client::PairedConnection`
pub trait FromResp: Sized {
    /// Return a `Result` containing either `Self` or `Error`.  Errors can occur due to either: a) the particular
    /// `RespValue` being incompatible with the required type, or b) a remote Redis error occuring.
    fn from_resp(resp: RespValue) -> Result<Self, Error> {
        Self::from_resp_int(resp.to_result()?)
    }

    fn from_resp_int(resp: RespValue) -> Result<Self, Error>;
}

impl FromResp for RespValue {
    fn from_resp_int(resp: RespValue) -> Result<RespValue, Error> {
        Ok(resp)
    }
}

impl FromResp for String {
    fn from_resp_int(resp: RespValue) -> Result<String, Error> {
        match resp {
            RespValue::BulkString(ref bytes) => Ok(String::from_utf8_lossy(bytes).into_owned()),
            RespValue::Integer(i) => Ok(i.to_string()),
            RespValue::SimpleString(string) => Ok(string),
            _ => Err(error::resp("Cannot convert into a string", resp)),
        }
    }
}

impl FromResp for usize {
    fn from_resp_int(resp: RespValue) -> Result<usize, Error> {
        match resp {
            RespValue::Error(string) => Err(Error::Remote(string)),
            RespValue::Integer(i) => Ok(i),
            _ => Err(error::resp("Cannot be converted into a usize", resp)),
        }
    }
}

impl FromResp for () {
    fn from_resp_int(resp: RespValue) -> Result<(), Error> {
        match resp {
            RespValue::Error(string) => Err(Error::Remote(string)),
            RespValue::SimpleString(string) => {
                match string.as_ref() {
                    "OK" => Ok(()),
                    _ => {
                        Err(Error::RESP(format!("Unexpected value within SimpleString: {}",
                                                string),
                                        None))
                    }
                }
            }
            _ => Err(error::resp("Unexpected value", resp)),
        }
    }
}

/// A trait to be implemented on types that can be automatically converted into `RespValue`s.
///
/// A `From<T>` where `T: ToResp` has been implemented so that everything that implements `ToResp`
/// can be converted with conversion traits.
pub trait ToResp {
    fn to_resp(&self) -> RespValue;
}

impl<'a, T> ToResp for &'a [T]
    where T: ToResp
{
    fn to_resp(&self) -> RespValue {
        RespValue::Array(self.as_ref().iter().map(|x| x.to_resp()).collect())
    }
}

impl<'a, T> ToResp for Vec<T>
    where T: ToResp
{
    fn to_resp(&self) -> RespValue {
        RespValue::Array(self.into_iter().map(|v| v.to_resp()).collect())
    }
}

impl<T: ToResp> From<T> for RespValue {
    fn from(from: T) -> RespValue {
        from.to_resp()
    }
}

/// A specific trait to convert into a `RespValue::BulkString`
pub trait ToRespString {
    fn to_resp_string(&self) -> RespValue;
}

impl ToRespString for String {
    fn to_resp_string(&self) -> RespValue {
        RespValue::BulkString(self.as_bytes().into())
    }
}

impl<'a> ToRespString for &'a str {
    fn to_resp_string(&self) -> RespValue {
        RespValue::BulkString(self.as_bytes().into())
    }
}

impl<T> ToResp for T
    where T: ToRespString
{
    fn to_resp(&self) -> RespValue {
        self.to_resp_string()
    }
}

/// Codec to read frames
pub struct RespCodec;

fn write_rn(buf: &mut BytesMut) {
    buf.put_u8(b'\r');
    buf.put_u8(b'\n');
}

fn check_and_reserve(buf: &mut BytesMut, amt: usize) {
    let remaining_bytes = buf.remaining_mut();
    if remaining_bytes < amt {
        buf.reserve(amt);
    }
}

fn write_header(symb: u8, len: usize, buf: &mut BytesMut) {
    let len_as_string = len.to_string();
    let len_as_bytes = len_as_string.as_bytes();
    let header_bytes = 1 + len_as_bytes.len() + 2;
    check_and_reserve(buf, header_bytes);
    buf.put_u8(symb);
    buf.extend(len_as_bytes);
    write_rn(buf);
}

fn write_simple_string(symb: u8, string: &str, buf: &mut BytesMut) {
    let bytes = string.as_bytes();
    let size = 1 + bytes.len() + 2;
    check_and_reserve(buf, size);
    buf.put_u8(symb);
    buf.extend(bytes);
    write_rn(buf);
}

impl Encoder for RespCodec {
    type Item = RespValue;
    type Error = io::Error;

    fn encode(&mut self, msg: RespValue, buf: &mut BytesMut) -> Result<(), Self::Error> {
        match msg {
            RespValue::Array(ary) => {
                write_header(b'*', ary.len(), buf);
                for v in ary {
                    self.encode(v, buf)?;
                }
            }
            RespValue::BulkString(bstr) => {
                let len = bstr.len();
                write_header(b'$', len, buf);
                check_and_reserve(buf, len + 2);
                buf.extend(bstr);
                write_rn(buf);
            }
            RespValue::Error(ref string) => {
                write_simple_string(b'-', string, buf);
            }
            RespValue::Integer(val) => {
                // Simple integer are just the header
                write_header(b':', val, buf);
            }
            RespValue::SimpleString(ref string) => {
                write_simple_string(b'+', string, buf);
            }
        }
        Ok(())
    }
}

#[inline]
fn parse_error(message: String) -> Error {
    Error::RESP(message, None)
}

/// Many RESP types have their length (which is either bytes or "number of elements", depending on context)
/// encoded as a string, terminated by "\r\n", this looks for them.
///
/// Only return the string if the whole sequence is complete, including the terminator bytes (but those final
/// two bytes will not be returned)
///
/// TODO - rename this function potentially, it's used for simple integers too
fn scan_integer<'a>(buf: &'a mut BytesMut, idx: usize) -> Result<Option<(usize, &'a [u8])>, Error> {
    let length = buf.len();
    let mut at_end = false;
    let mut pos = idx;
    loop {
        if length <= pos {
            return Ok(None);
        }
        match (at_end, buf[pos]) {
            (true, b'\n') => return Ok(Some((pos + 1, &buf[idx..pos - 1]))),
            (false, b'\r') => at_end = true,
            (false, b'0'...b'9') => (),
            (_, val) => return Err(parse_error(format!("Unexpected byte in size_string: {}", val))),
        }
        pos += 1;
    }
}

fn scan_string(buf: &mut BytesMut, idx: usize) -> Option<(usize, String)> {
    let length = buf.len();
    let mut at_end = false;
    let mut pos = idx;
    loop {
        if length <= pos {
            return None;
        }
        match (at_end, buf[pos]) {
            (true, b'\n') => {
                let value = String::from_utf8_lossy(&buf[idx..pos - 1]).into_owned();
                return Some((pos + 1, value));
            }
            (true, _) => at_end = false,
            (false, b'\r') => at_end = true,
            (false, _) => (),
        }
        pos += 1;
    }
}

fn decode_raw_integer(buf: &mut BytesMut, idx: usize) -> Result<Option<(usize, usize)>, Error> {
    match scan_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, int_str))) => {
            let int: usize = str::from_utf8(int_str)
                .expect("Not a string")
                .parse()
                .expect("Not an integer");
            Ok(Some((pos, int)))
        }
        Err(e) => Err(e),
    }
}

type DecodeResult = Result<Option<(usize, RespValue)>, Error>;

fn decode_bulk_string(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match decode_raw_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, size))) => {
            let remaining = buf.len() - pos;
            let required_bytes = size + 2;

            if remaining < required_bytes {
                return Ok(None);
            }

            let bulk_string = RespValue::BulkString(buf[pos..(pos + size)].to_vec());
            Ok(Some((pos + required_bytes, bulk_string)))
        }
        Err(e) => Err(e),
    }
}

fn decode_array(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match decode_raw_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, size))) => {
            let mut pos = pos;
            let mut values = Vec::with_capacity(size);
            for _ in 0..size {
                match decode(buf, pos) {
                    Ok(None) => return Ok(None),
                    Ok(Some((new_pos, value))) => {
                        values.push(value);
                        pos = new_pos;
                    }
                    Err(e) => return Err(e),
                }
            }
            Ok(Some((pos, RespValue::Array(values))))
        }
        Err(e) => Err(e),
    }
}

fn decode_integer(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match decode_raw_integer(buf, idx) {
        Ok(None) => Ok(None),
        Ok(Some((pos, int))) => Ok(Some((pos, RespValue::Integer(int)))),
        Err(e) => Err(e),
    }
}

/// A simple string is any series of bytes that ends with `\r\n`
fn decode_simple_string(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match scan_string(buf, idx) {
        None => Ok(None),
        Some((pos, string)) => Ok(Some((pos, RespValue::SimpleString(string)))),
    }
}

fn decode_error(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    match scan_string(buf, idx) {
        None => Ok(None),
        Some((pos, string)) => Ok(Some((pos, RespValue::Error(string)))),
    }
}

fn decode(buf: &mut BytesMut, idx: usize) -> DecodeResult {
    let length = buf.len();
    if length <= idx {
        return Ok(None);
    }

    let first_byte = buf[idx];
    match first_byte {
        b'$' => decode_bulk_string(buf, idx + 1),
        b'*' => decode_array(buf, idx + 1),
        b':' => decode_integer(buf, idx + 1),
        b'+' => decode_simple_string(buf, idx + 1),
        b'-' => decode_error(buf, idx + 1),
        _ => Err(parse_error(format!("Unexpected byte: {}", first_byte))),
    }
}

impl Decoder for RespCodec {
    type Item = RespValue;
    type Error = Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        match decode(buf, 0) {
            Ok(None) => Ok(None),
            Ok(Some((pos, item))) => {
                buf.split_to(pos);
                Ok(Some(item))
            }
            Err(e) => Err(e),
        }
    }
}

#[cfg(test)]
mod tests {
    use bytes::BytesMut;

    use tokio_io::codec::{Decoder, Encoder};

    use super::{RespCodec, RespValue};

    #[test]
    fn test_bulk_string() {
        let resp_object = RespValue::BulkString("THISISATEST".as_bytes().to_vec());
        let mut bytes = BytesMut::new();
        let mut codec = RespCodec;
        codec.encode(resp_object.clone(), &mut bytes).unwrap();
        assert_eq!(b"$11\r\nTHISISATEST\r\n".to_vec(), bytes.to_vec());

        let deserialized = codec.decode(&mut bytes).unwrap().unwrap();
        assert_eq!(deserialized, resp_object);
    }

    #[test]
    fn test_array() {
        let resp_object = RespValue::Array(vec!["TEST1".into(), "TEST2".into()]);
        let mut bytes = BytesMut::new();
        let mut codec = RespCodec;
        codec.encode(resp_object.clone(), &mut bytes).unwrap();
        assert_eq!(b"*2\r\n$5\r\nTEST1\r\n$5\r\nTEST2\r\n".to_vec(),
                   bytes.to_vec());

        let deserialized = codec.decode(&mut bytes).unwrap().unwrap();
        assert_eq!(deserialized, resp_object);
    }
}<|endoftext|>


<commit_before>//! Functions relating to network operations.
//!
//! This will eventually be moved

pub mod deque_buffer;

pub mod packet;

use std::collections::VecDeque;

use std::net::{Shutdown, TcpListener, TcpStream};

use std::thread;

use std::io::Read;

use self::deque_buffer::DequeBuffer;

fn start_listening(port: u16) -> Result<TcpListener, &'static str> {

    let listener = TcpListener::bind(("127.0.0.1", port));

    if listener.is_ok() {
        return Ok(listener.unwrap());
    } else {
        return Err("Could not listen on requested port");
    }

}

pub fn start_network() {

    let listener = start_listening(25565);

    if listener.is_err() {
        error!("Could not start listening on port {}", 25565);
    }

    let listener = listener.unwrap();

    for stream in listener.incoming() {

        match stream {
            Ok(stream) => {
                thread::spawn(move || {
                    // New connection succeeded
                    handle_new_client(stream)
                });
            }
            Err(e) => {
                error!("Could not accept a new stream: {}", e);
            }
        }

    }

    drop(listener);

}

pub fn handle_new_client(mut stream: TcpStream) {
    info!("Got a new client from {}! Oh my!", stream.peer_addr().unwrap());

    loop {

        let mut raw_buffer = vec![0u8; 512];

        match stream.read(&mut raw_buffer) {

            Err(error) => {
                // Connection dropped out or something broke
                // In the future, we use the customized NetherrackStream to see if we're in GAME mode, and if so, get the attached player and start saving their information and get rid of them
                // For now, we just exit the loop
                error!("Got an error! {}", error);

                let _ = stream.shutdown(Shutdown::Both);
                break; //Close the connection
            }
            Ok(count) => {
                if count == 0 {
                    //Nothing in stream, let's sleep for a few milliseconds and try again
                    ::std::thread::sleep_ms(2);
                } else {

                    // Now truncate the raw buffer - Don't worry, it will reallocate to 512 on the next networking loop
                    raw_buffer.truncate(count);

                    let mut backing: VecDeque<u8> = VecDeque::<u8>::with_capacity(count);

                    for b in raw_buffer {

                        backing.push_back(b);
                    }

                    let mut buffer = DequeBuffer::from(backing);

                    //drop(raw_buffer);

                    let length_result = buffer.read_unsigned_varint_32();
                    let id_result = buffer.read_unsigned_varint_32();

                    if length_result.is_err() || id_result.is_err() {
                        error!("Error while reading packet ID or length! Oh no, dropping connection");

                        debug!("Shutdown result: {:?}", stream.shutdown(Shutdown::Both));
                        break;
                    }

                    info!("Read {} bytes, length is {}, id is {}", count, length_result.unwrap(), id_result.unwrap());

                    info!("Remaining buffer size is {}", buffer.remaining());

                    // This is a memory leak. We need to get the time of last packet and check to see if it's more than 20 seconds ago
                }
            }

        }

    }

    info!("Connection closed");

}<commit_msg>Time out after 10 seconds without a valid packet<commit_after>//! Functions relating to network operations.
//!
//! This will eventually be moved

extern crate chrono;
use self::chrono::duration::Duration;

pub mod deque_buffer;

pub mod packet;
use self::packet::PacketHeader;

use std::collections::VecDeque;

use std::net::{Shutdown, TcpListener, TcpStream};

use std::thread;

use std::io::Read;

use self::deque_buffer::DequeBuffer;

fn start_listening(port: u16) -> Result<TcpListener, &'static str> {

    let listener = TcpListener::bind(("127.0.0.1", port));

    if listener.is_ok() {
        return Ok(listener.unwrap());
    } else {
        return Err("Could not listen on requested port");
    }

}

pub fn start_network() {

    let listener = start_listening(25565);

    if listener.is_err() {
        error!("Could not start listening on port {}", 25565);
    }

    let listener = listener.unwrap();

    for stream in listener.incoming() {

        match stream {
            Ok(stream) => {
                thread::spawn(move || {
                    // New connection succeeded
                    handle_new_client(stream)
                });
            }
            Err(e) => {
                error!("Could not accept a new stream: {}", e);
            }
        }

    }

    drop(listener);

}

pub fn handle_new_client(mut stream: TcpStream) {
    info!("Got a new client from {}! Oh my!", stream.peer_addr().unwrap());

    let mut timeout_duration = Duration::zero();

    loop {

        if (timeout_duration.num_seconds() >= 10) {
            info!("Client timed out; Shutdown was {:?}", stream.shutdown(Shutdown::Both));
            break;
        }

        let mut raw_buffer = vec![0u8; 512];

        match stream.read(&mut raw_buffer) {

            Err(error) => {
                // Connection dropped out or something broke
                // In the future, we use the customized NetherrackStream to see if we're in GAME mode, and if so, get the attached player and start saving their information and get rid of them
                // For now, we just exit the loop
                error!("Got an error! {}", error);

                let _ = stream.shutdown(Shutdown::Both);
                break; //Close the connection
            }
            Ok(count) => {
                if count == 0 {
                    //Nothing in stream, let's sleep for a few milliseconds and try again
                    timeout_duration = timeout_duration + Duration::milliseconds(2);
                    ::std::thread::sleep_ms(2);
                } else {

                    timeout_duration = Duration::zero();

                    // Now truncate the raw buffer - Don't worry, it will reallocate to 512 on the next networking loop
                    raw_buffer.truncate(count);

                    let mut backing: VecDeque<u8> = VecDeque::<u8>::with_capacity(count);

                    for b in raw_buffer {

                        backing.push_back(b);
                    }

                    let mut buffer = DequeBuffer::from(backing);

                    //drop(raw_buffer);

                    let length_result = buffer.read_unsigned_varint_32();
                    let id_result = buffer.read_unsigned_varint_32();

                    if length_result.is_err() || id_result.is_err() {
                        error!("Error while reading packet ID or length! Oh no, dropping connection");

                        debug!("Shutdown result: {:?}", stream.shutdown(Shutdown::Both));
                        break;
                    }

                    let packet_header = PacketHeader { length: length_result.unwrap(), id: id_result.unwrap() };

                    info!("Read {} bytes, length is {}, id is {}", count, packet_header.length, packet_header.id);

                    info!("Remaining buffer size is {}", buffer.remaining());

                    //TODO: Move to a MinecraftClientConnection, determine if the client is Forge or Vanilla, and set the GameState to HANDSHAKE

                    if packet_header.id == 0 {
                        if packet_header.length == 1 {

                            //This is a request packet
                            debug!("Received a Server Status Request packet");

                        } else {

                            //This is a handshake packet
                            debug!("Received a Handshake packet");

                            let protocol_result = buffer.read_unsigned_varint_32();

                            if protocol_result.is_err() {
                                error!("Error reading protocol version");
                                debug!("Shutdown result: {:?}", stream.shutdown(Shutdown::Both));
                                break;
                            }

                            debug!("Protocol version {}", protocol_result.unwrap());

                        }
                    }

                    // This is a memory leak. We need to get the time of last packet and check to see if it's more than 20 seconds ago
                }
            }

        }

    }

    info!("Connection closed");

}<|endoftext|>


<commit_before>use super::*;
use crate::norm::Norm;
use num_traits::One;
use std::iter::*;

/// Execute Arnodi iteration as Rust iterator
///
/// - [Arnoldi iteration - Wikipedia](https://en.wikipedia.org/wiki/Arnoldi_iteration)
///
pub struct Arnoldi<A, S, F, Ortho>
where
    A: Scalar,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    a: F,
    /// Next vector (normalized `|v|=1`)
    v: ArrayBase<S, Ix1>,
    /// Orthogonalizer
    ortho: Ortho,
    /// Coefficients to be composed into H-matrix
    h: Vec<Array1<A>>,
}

impl<A, S, F, Ortho> Arnoldi<A, S, F, Ortho>
where
    A: Scalar + Lapack,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    pub fn new(a: F, mut v: ArrayBase<S, Ix1>, mut ortho: Ortho) -> Self {
        assert_eq!(ortho.len(), 0);
        assert!(ortho.tolerance() < One::one());
        // normalize before append because |v| may be smaller than ortho.tolerance()
        let norm = v.norm_l2();
        azip!(mut v(&mut v) in { *v = v.div_real(norm) });
        ortho.append(v.view());
        Arnoldi {
            a,
            v,
            ortho,
            h: Vec::new(),
        }
    }

    /// Dimension of Krylov subspace
    pub fn dim(&self) -> usize {
        self.ortho.len()
    }

    /// Iterate until convergent
    pub fn complete(mut self) -> (Q<A>, H<A>) {
        for _ in &mut self {} // execute iteration until convergent
        let q = self.ortho.get_q();
        let n = self.dim();
        let mut h = Array2::zeros((n, n).f());
        for (i, hc) in self.h.iter().enumerate() {
            let m = std::cmp::max(n, i + 1);
            for j in 0..m {
                h[(j, i)] = hc[j];
            }
        }
        (q, h)
    }
}

impl<A, S, F, Ortho> Iterator for Arnoldi<A, S, F, Ortho>
where
    A: Scalar,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    type Item = Array1<A>;

    fn next(&mut self) -> Option<Self::Item> {
        (self.a)(&mut self.v);
        let result = self.ortho.div_append(&mut self.v);
        azip!(mut v(&mut self.v) in { *v = v.div_real(result.residual_norm()) });
        if result.is_dependent() {
            None
        } else {
            let coef = result.into_coeff();
            self.h.push(coef.clone());
            Some(coef)
        }
    }
}

/// Interpret a matrix as a linear operator
pub fn mul_mat<A, S1, S2>(a: ArrayBase<S1, Ix2>) -> impl Fn(&mut ArrayBase<S2, Ix1>)
where
    A: Scalar,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    let (n, m) = a.dim();
    assert_eq!(n, m, "Input matrix must be square");
    move |x| {
        assert_eq!(m, x.len(), "Input matrix and vector sizes mismatch");
        let ax = a.dot(x);
        azip!(mut x(x), ax in { *x = ax });
    }
}

/// Utility to execute Arnoldi iteration with Householder reflection
pub fn arnoldi_householder<A, S1, S2>(a: ArrayBase<S1, Ix2>, v: ArrayBase<S2, Ix1>, tol: A::Real) -> (Q<A>, H<A>)
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    let householder = Householder::new(v.len(), tol);
    Arnoldi::new(mul_mat(a), v, householder).complete()
}

/// Utility to execute Arnoldi iteration with modified Gram-Schmit orthogonalizer
pub fn arnoldi_mgs<A, S1, S2>(a: ArrayBase<S1, Ix2>, v: ArrayBase<S2, Ix1>, tol: A::Real) -> (Q<A>, H<A>)
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    let mgs = MGS::new(v.len(), tol);
    Arnoldi::new(mul_mat(a), v, mgs).complete()
}
<commit_msg>Add test<commit_after>use super::*;
use crate::norm::Norm;
use num_traits::One;
use std::iter::*;

/// Execute Arnodi iteration as Rust iterator
///
/// - [Arnoldi iteration - Wikipedia](https://en.wikipedia.org/wiki/Arnoldi_iteration)
///
pub struct Arnoldi<A, S, F, Ortho>
where
    A: Scalar,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    a: F,
    /// Next vector (normalized `|v|=1`)
    v: ArrayBase<S, Ix1>,
    /// Orthogonalizer
    ortho: Ortho,
    /// Coefficients to be composed into H-matrix
    h: Vec<Array1<A>>,
}

impl<A, S, F, Ortho> Arnoldi<A, S, F, Ortho>
where
    A: Scalar + Lapack,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    pub fn new(a: F, mut v: ArrayBase<S, Ix1>, mut ortho: Ortho) -> Self {
        assert_eq!(ortho.len(), 0);
        assert!(ortho.tolerance() < One::one());
        // normalize before append because |v| may be smaller than ortho.tolerance()
        let norm = v.norm_l2();
        azip!(mut v(&mut v) in { *v = v.div_real(norm) });
        ortho.append(v.view());
        Arnoldi {
            a,
            v,
            ortho,
            h: Vec::new(),
        }
    }

    /// Dimension of Krylov subspace
    pub fn dim(&self) -> usize {
        self.ortho.len()
    }

    /// Iterate until convergent
    pub fn complete(mut self) -> (Q<A>, H<A>) {
        for _ in &mut self {} // execute iteration until convergent
        let q = self.ortho.get_q();
        let n = self.h.len();
        let mut h = Array2::zeros((n, n).f());
        for (i, hc) in self.h.iter().enumerate() {
            let m = std::cmp::min(n, i + 2);
            for j in 0..m {
                h[(j, i)] = hc[j];
            }
        }
        (q, h)
    }
}

impl<A, S, F, Ortho> Iterator for Arnoldi<A, S, F, Ortho>
where
    A: Scalar,
    S: DataMut<Elem = A>,
    F: Fn(&mut ArrayBase<S, Ix1>),
    Ortho: Orthogonalizer<Elem = A>,
{
    type Item = Array1<A>;

    fn next(&mut self) -> Option<Self::Item> {
        (self.a)(&mut self.v);
        let result = self.ortho.div_append(&mut self.v);
        azip!(mut v(&mut self.v) in { *v = v.div_real(result.residual_norm()) });
        match result {
            AppendResult::Added(coef) => {
                dbg!(&coef);
                self.h.push(coef.clone());
                Some(coef)
            }
            AppendResult::Dependent(coef) => {
                dbg!(&coef);
                self.h.push(coef);
                None
            }
        }
    }
}

/// Interpret a matrix as a linear operator
pub fn mul_mat<A, S1, S2>(a: ArrayBase<S1, Ix2>) -> impl Fn(&mut ArrayBase<S2, Ix1>)
where
    A: Scalar,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    let (n, m) = a.dim();
    assert_eq!(n, m, "Input matrix must be square");
    move |x| {
        assert_eq!(m, x.len(), "Input matrix and vector sizes mismatch");
        let ax = a.dot(x);
        azip!(mut x(x), ax in { *x = ax });
    }
}

/// Utility to execute Arnoldi iteration with Householder reflection
pub fn arnoldi_householder<A, S1, S2>(a: ArrayBase<S1, Ix2>, v: ArrayBase<S2, Ix1>, tol: A::Real) -> (Q<A>, H<A>)
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    let householder = Householder::new(v.len(), tol);
    Arnoldi::new(mul_mat(a), v, householder).complete()
}

/// Utility to execute Arnoldi iteration with modified Gram-Schmit orthogonalizer
pub fn arnoldi_mgs<A, S1, S2>(a: ArrayBase<S1, Ix2>, v: ArrayBase<S2, Ix1>, tol: A::Real) -> (Q<A>, H<A>)
where
    A: Scalar + Lapack,
    S1: Data<Elem = A>,
    S2: DataMut<Elem = A>,
{
    let mgs = MGS::new(v.len(), tol);
    Arnoldi::new(mul_mat(a), v, mgs).complete()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::generate::*;

    #[test]
    fn aq_qh() {
        let a = array![[1.0, 2.0], [3.0, 5.0]];
        let mut v = Array::zeros(2);
        v[0] = 1.0;
        let (q, h) = arnoldi_mgs(a.clone(), v, 1e-9);
        println!("A = \n{:?}", &a);
        println!("Q = \n{:?}", &q);
        println!("H = \n{:?}", &h);
        println!("AQ = \n{:?}", a.dot(&q));
        println!("QH = \n{:?}", q.dot(&h));
        panic!()
    }

    #[test]
    fn aq_qh_random() {
        let a: Array2<f64> = random((5, 5));
        let mut v = Array::zeros(5);
        v[0] = 1.0;
        let (q, h) = arnoldi_mgs(a.clone(), v, 1e-9);
        println!("A = \n{:?}", &a);
        println!("Q = \n{:?}", &q);
        println!("H = \n{:?}", &h);
        println!("AQ = \n{:?}", a.dot(&q));
        println!("QH = \n{:?}", q.dot(&h));
        panic!()
    }

}
<|endoftext|>


<commit_before>#![cfg(test)]

extern crate compiletest_rs as compiletest;

use std::path::PathBuf;

#[test]
#[cfg(feature = "unstable")]
fn negative_tests_compile_fail_unstable() {
    let mode = "compile-fail";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/compile-fail-unstable");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
fn negative_tests_run_fail() {
    let mode = "run-fail";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/run-fail");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
#[cfg(feature = "unstable")]
fn negative_tests_run_fail_unstable() {
    let mode = "run-fail";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/run-fail-unstable");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
fn positive_test_run_pass() {
    let mode = "run-pass";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/run-pass");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
#[cfg(feature = "unstable")]
fn positive_test_run_pass_unstable() {
    let mode = "run-pass";
    let mut config = compiletest::default_config();
    let cfg_mode = mode.parse().ok().expect("Invalid mode");

    config.mode = cfg_mode;
    config.src_base = PathBuf::from("tests/run-pass-unstable");
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}
<commit_msg>Factor out common code in compile tests<commit_after>#![cfg(test)]

extern crate compiletest_rs as compiletest;

use std::path::PathBuf;

fn run_compiletest(mode: &str, path: &str) {
    let mut config = compiletest::default_config();
    config.mode = mode.parse().ok().expect("Invalid mode");
    config.src_base = PathBuf::from(path);
    config.target_rustcflags = Some("-L target/debug/ -L target/debug/deps/".to_owned());

    compiletest::run_tests(&config);
}

#[test]
#[cfg(feature = "unstable")]
fn negative_tests_compile_fail_unstable() {
    run_compiletest("compile-fail", "tests/compile-fail-unstable");
}

#[test]
fn negative_tests_run_fail() {
    run_compiletest("run-fail", "tests/run-fail");
}

#[test]
#[cfg(feature = "unstable")]
fn negative_tests_run_fail_unstable() {
    run_compiletest("run-fail", "tests/run-fail-unstable");
}

#[test]
fn positive_test_run_pass() {
    run_compiletest("run-pass", "tests/run-pass");
}

#[test]
#[cfg(feature = "unstable")]
fn positive_test_run_pass_unstable() {
    run_compiletest("run-pass", "tests/run-pass-unstable");
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

// The azure unit tests require this
// on Linux to link correctly.

#[cfg(target_os = "linux")]
#[link(name = "GL")]
extern { }

#[test]
fn sanity_check() {
    use azure::AzSanityCheck;

    unsafe { AzSanityCheck() };
}
<commit_msg>Add a test-only dependency on skia so that cargo test works on this crate.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

// The azure unit tests require this
// on Linux to link correctly.

#[cfg(target_os = "linux")]
#[link(name = "skia")]
#[link(name = "GL")]
extern { }

#[test]
fn sanity_check() {
    use azure::AzSanityCheck;

    unsafe { AzSanityCheck() };
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::mpsc::channel;
use std::thread;

use kvproto::kvrpcpb::Context;

use test_raftstore::*;
use tikv::raftstore::store::engine::IterOption;
use tikv::storage::engine::*;
use tikv::storage::{CFStatistics, CfName, Key, CF_DEFAULT};
use tikv::util::codec::bytes;
use tikv::util::escape;
use tikv::util::HandyRwLock;

#[test]
fn test_raftkv() {
    let count = 1;
    let mut cluster = new_server_cluster(0, count);
    cluster.run();

    // make sure leader has been elected.
    assert_eq!(cluster.must_get(b"k1"), None);

    let region = cluster.get_region(b"");
    let leader_id = cluster.leader_of_region(region.get_id()).unwrap();
    let storage = cluster.sim.rl().storages[&leader_id.get_id()].clone();

    let mut ctx = Context::new();
    ctx.set_region_id(region.get_id());
    ctx.set_region_epoch(region.get_region_epoch().clone());
    ctx.set_peer(region.get_peers()[0].clone());

    get_put(&ctx, &storage);
    batch(&ctx, &storage);
    seek(&ctx, &storage);
    near_seek(&ctx, &storage);
    cf(&ctx, &storage);
    empty_write(&ctx, &storage);
    wrong_context(&ctx, &storage);
    empty_batch_snapshot(&storage);
    // TODO: test multiple node
}

#[test]
fn test_read_leader_in_lease() {
    let count = 3;
    let mut cluster = new_server_cluster(0, count);
    cluster.run();

    let k1 = b"k1";
    let (k2, v2) = (b"k2", b"v2");

    // make sure leader has been elected.
    assert_eq!(cluster.must_get(k1), None);

    let region = cluster.get_region(b"");
    let leader = cluster.leader_of_region(region.get_id()).unwrap();
    let storage = cluster.sim.rl().storages[&leader.get_id()].clone();

    let mut ctx = Context::new();
    ctx.set_region_id(region.get_id());
    ctx.set_region_epoch(region.get_region_epoch().clone());
    ctx.set_peer(leader.clone());

    // write some data
    assert_none(&ctx, &storage, k2);
    must_put(&ctx, &storage, k2, v2);

    // isolate leader
    cluster.add_send_filter(IsolationFilterFactory::new(leader.get_store_id()));

    // leader still in lease, check if can read on leader
    assert_eq!(can_read(&ctx, &storage, k2, v2), true);
}

#[test]
fn test_batch_snapshot() {
    let count = 3;
    let mut cluster = new_server_cluster(0, count);
    cluster.run();

    let key = b"key";
    // make sure leader has been elected.
    assert_eq!(cluster.must_get(key), None);

    let region = cluster.get_region(b"");
    let leader = cluster.leader_of_region(region.get_id()).unwrap();
    let storage = cluster.sim.rl().storages[&leader.get_id()].clone();

    let mut ctx = Context::new();
    ctx.set_region_id(region.get_id());
    ctx.set_region_epoch(region.get_region_epoch().clone());
    ctx.set_peer(leader.clone());

    let size = 3;
    let batch = vec![ctx.clone(); size];
    let snapshots = must_batch_snapshot(batch, &storage);
    assert_eq!(size, snapshots.len());
    for s in snapshots {
        assert!(s.is_some());
    }
    // sleep util leader lease is expired.
    thread::sleep(cluster.cfg.raft_store.raft_store_max_leader_lease.0);
    let batch = vec![ctx; size];
    let snapshots = must_batch_snapshot(batch, &storage);
    assert_eq!(size, snapshots.len());
    for s in snapshots {
        assert!(s.is_none());
    }
}

fn must_batch_snapshot<E: Engine>(batch: Vec<Context>, engine: &E) -> BatchResults<E::Snap> {
    let (tx, rx) = channel();
    let on_finished = box move |snapshots| {
        tx.send(snapshots).unwrap();
    };
    engine.async_batch_snapshot(batch, on_finished).unwrap();
    rx.recv().unwrap()
}

fn must_put<E: Engine>(ctx: &Context, engine: &E, key: &[u8], value: &[u8]) {
    engine.put(ctx, Key::from_raw(key), value.to_vec()).unwrap();
}

fn must_put_cf<E: Engine>(ctx: &Context, engine: &E, cf: CfName, key: &[u8], value: &[u8]) {
    engine
        .put_cf(ctx, cf, Key::from_raw(key), value.to_vec())
        .unwrap();
}

fn must_delete<E: Engine>(ctx: &Context, engine: &E, key: &[u8]) {
    engine.delete(ctx, Key::from_raw(key)).unwrap();
}

fn must_delete_cf<E: Engine>(ctx: &Context, engine: &E, cf: CfName, key: &[u8]) {
    engine.delete_cf(ctx, cf, Key::from_raw(key)).unwrap();
}

fn assert_has<E: Engine>(ctx: &Context, engine: &E, key: &[u8], value: &[u8]) {
    let snapshot = engine.snapshot(ctx).unwrap();
    assert_eq!(snapshot.get(&Key::from_raw(key)).unwrap().unwrap(), value);
}

fn can_read<E: Engine>(ctx: &Context, engine: &E, key: &[u8], value: &[u8]) -> bool {
    if let Ok(s) = engine.snapshot(ctx) {
        assert_eq!(s.get(&Key::from_raw(key)).unwrap().unwrap(), value);
        return true;
    }
    false
}

fn assert_has_cf<E: Engine>(ctx: &Context, engine: &E, cf: CfName, key: &[u8], value: &[u8]) {
    let snapshot = engine.snapshot(ctx).unwrap();
    assert_eq!(
        snapshot.get_cf(cf, &Key::from_raw(key)).unwrap().unwrap(),
        value
    );
}

fn assert_none<E: Engine>(ctx: &Context, engine: &E, key: &[u8]) {
    let snapshot = engine.snapshot(ctx).unwrap();
    assert_eq!(snapshot.get(&Key::from_raw(key)).unwrap(), None);
}

fn assert_none_cf<E: Engine>(ctx: &Context, engine: &E, cf: CfName, key: &[u8]) {
    let snapshot = engine.snapshot(ctx).unwrap();
    assert_eq!(snapshot.get_cf(cf, &Key::from_raw(key)).unwrap(), None);
}

fn assert_seek<E: Engine>(ctx: &Context, engine: &E, key: &[u8], pair: (&[u8], &[u8])) {
    let snapshot = engine.snapshot(ctx).unwrap();
    let mut cursor = snapshot
        .iter(IterOption::default(), ScanMode::Mixed)
        .unwrap();
    let mut statistics = CFStatistics::default();
    cursor.seek(&Key::from_raw(key), &mut statistics).unwrap();
    assert_eq!(cursor.key(&mut statistics), &*bytes::encode_bytes(pair.0));
    assert_eq!(cursor.value(&mut statistics), pair.1);
}

fn assert_seek_cf<E: Engine>(
    ctx: &Context,
    engine: &E,
    cf: CfName,
    key: &[u8],
    pair: (&[u8], &[u8]),
) {
    let snapshot = engine.snapshot(ctx).unwrap();
    let mut cursor = snapshot
        .iter_cf(cf, IterOption::default(), ScanMode::Mixed)
        .unwrap();
    let mut statistics = CFStatistics::default();
    cursor.seek(&Key::from_raw(key), &mut statistics).unwrap();
    assert_eq!(cursor.key(&mut statistics), &*bytes::encode_bytes(pair.0));
    assert_eq!(cursor.value(&mut statistics), pair.1);
}

fn assert_near_seek<I: Iterator>(cursor: &mut Cursor<I>, key: &[u8], pair: (&[u8], &[u8])) {
    let mut statistics = CFStatistics::default();
    assert!(
        cursor
            .near_seek(&Key::from_raw(key), &mut statistics)
            .unwrap(),
        escape(key)
    );
    assert_eq!(cursor.key(&mut statistics), &*bytes::encode_bytes(pair.0));
    assert_eq!(cursor.value(&mut statistics), pair.1);
}

fn assert_near_reverse_seek<I: Iterator>(cursor: &mut Cursor<I>, key: &[u8], pair: (&[u8], &[u8])) {
    let mut statistics = CFStatistics::default();
    assert!(
        cursor
            .near_reverse_seek(&Key::from_raw(key), &mut statistics)
            .unwrap(),
        escape(key)
    );
    assert_eq!(cursor.key(&mut statistics), &*bytes::encode_bytes(pair.0));
    assert_eq!(cursor.value(&mut statistics), pair.1);
}

fn get_put<E: Engine>(ctx: &Context, engine: &E) {
    assert_none(ctx, engine, b"x");
    must_put(ctx, engine, b"x", b"1");
    assert_has(ctx, engine, b"x", b"1");
    must_put(ctx, engine, b"x", b"2");
    assert_has(ctx, engine, b"x", b"2");
}

fn batch<E: Engine>(ctx: &Context, engine: &E) {
    engine
        .write(
            ctx,
            vec![
                Modify::Put(CF_DEFAULT, Key::from_raw(b"x"), b"1".to_vec()),
                Modify::Put(CF_DEFAULT, Key::from_raw(b"y"), b"2".to_vec()),
            ],
        )
        .unwrap();
    assert_has(ctx, engine, b"x", b"1");
    assert_has(ctx, engine, b"y", b"2");

    engine
        .write(
            ctx,
            vec![
                Modify::Delete(CF_DEFAULT, Key::from_raw(b"x")),
                Modify::Delete(CF_DEFAULT, Key::from_raw(b"y")),
            ],
        )
        .unwrap();
    assert_none(ctx, engine, b"y");
    assert_none(ctx, engine, b"y");
}

fn seek<E: Engine>(ctx: &Context, engine: &E) {
    must_put(ctx, engine, b"x", b"1");
    assert_seek(ctx, engine, b"x", (b"x", b"1"));
    assert_seek(ctx, engine, b"a", (b"x", b"1"));
    must_put(ctx, engine, b"z", b"2");
    assert_seek(ctx, engine, b"y", (b"z", b"2"));
    assert_seek(ctx, engine, b"x\x00", (b"z", b"2"));
    let snapshot = engine.snapshot(ctx).unwrap();
    let mut iter = snapshot
        .iter(IterOption::default(), ScanMode::Mixed)
        .unwrap();
    let mut statistics = CFStatistics::default();
    assert!(
        !iter
            .seek(&Key::from_raw(b"z\x00"), &mut statistics)
            .unwrap()
    );
    must_delete(ctx, engine, b"x");
    must_delete(ctx, engine, b"z");
}

fn near_seek<E: Engine>(ctx: &Context, engine: &E) {
    must_put(ctx, engine, b"x", b"1");
    must_put(ctx, engine, b"z", b"2");
    let snapshot = engine.snapshot(ctx).unwrap();
    let mut cursor = snapshot
        .iter(IterOption::default(), ScanMode::Mixed)
        .unwrap();
    assert_near_seek(&mut cursor, b"x", (b"x", b"1"));
    assert_near_seek(&mut cursor, b"a", (b"x", b"1"));
    assert_near_reverse_seek(&mut cursor, b"z1", (b"z", b"2"));
    assert_near_reverse_seek(&mut cursor, b"x1", (b"x", b"1"));
    assert_near_seek(&mut cursor, b"y", (b"z", b"2"));
    assert_near_seek(&mut cursor, b"x\x00", (b"z", b"2"));
    let mut statistics = CFStatistics::default();
    assert!(
        !cursor
            .near_seek(&Key::from_raw(b"z\x00"), &mut statistics)
            .unwrap()
    );
    must_delete(ctx, engine, b"x");
    must_delete(ctx, engine, b"z");
}

fn cf<E: Engine>(ctx: &Context, engine: &E) {
    assert_none_cf(ctx, engine, "default", b"key");
    must_put_cf(ctx, engine, "default", b"key", b"value");
    assert_has_cf(ctx, engine, "default", b"key", b"value");
    assert_seek_cf(ctx, engine, "default", b"k", (b"key", b"value"));
    must_delete_cf(ctx, engine, "default", b"key");
    assert_none_cf(ctx, engine, "default", b"key");
}

fn empty_write<E: Engine>(ctx: &Context, engine: &E) {
    engine.write(ctx, vec![]).unwrap_err();
}

fn wrong_context<E: Engine>(ctx: &Context, engine: &E) {
    let region_id = ctx.get_region_id();
    let mut ctx = ctx.to_owned();
    ctx.set_region_id(region_id + 1);
    assert!(engine.write(&ctx, vec![]).is_err());
}

fn empty_batch_snapshot<E: Engine>(engine: &E) {
    let on_finished = box move |_| {};
    engine
        .async_batch_snapshot(vec![], on_finished)
        .unwrap_err();
}
<commit_msg>tests: stabilize test_batch_snapshot (#3544)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::mpsc::channel;
use std::thread;

use kvproto::kvrpcpb::Context;

use test_raftstore::*;
use tikv::raftstore::store::engine::IterOption;
use tikv::storage::engine::*;
use tikv::storage::{CFStatistics, CfName, Key, CF_DEFAULT};
use tikv::util::codec::bytes;
use tikv::util::escape;
use tikv::util::HandyRwLock;

#[test]
fn test_raftkv() {
    let count = 1;
    let mut cluster = new_server_cluster(0, count);
    cluster.run();

    // make sure leader has been elected.
    assert_eq!(cluster.must_get(b"k1"), None);

    let region = cluster.get_region(b"");
    let leader_id = cluster.leader_of_region(region.get_id()).unwrap();
    let storage = cluster.sim.rl().storages[&leader_id.get_id()].clone();

    let mut ctx = Context::new();
    ctx.set_region_id(region.get_id());
    ctx.set_region_epoch(region.get_region_epoch().clone());
    ctx.set_peer(region.get_peers()[0].clone());

    get_put(&ctx, &storage);
    batch(&ctx, &storage);
    seek(&ctx, &storage);
    near_seek(&ctx, &storage);
    cf(&ctx, &storage);
    empty_write(&ctx, &storage);
    wrong_context(&ctx, &storage);
    empty_batch_snapshot(&storage);
    // TODO: test multiple node
}

#[test]
fn test_read_leader_in_lease() {
    let count = 3;
    let mut cluster = new_server_cluster(0, count);
    cluster.run();

    let k1 = b"k1";
    let (k2, v2) = (b"k2", b"v2");

    // make sure leader has been elected.
    assert_eq!(cluster.must_get(k1), None);

    let region = cluster.get_region(b"");
    let leader = cluster.leader_of_region(region.get_id()).unwrap();
    let storage = cluster.sim.rl().storages[&leader.get_id()].clone();

    let mut ctx = Context::new();
    ctx.set_region_id(region.get_id());
    ctx.set_region_epoch(region.get_region_epoch().clone());
    ctx.set_peer(leader.clone());

    // write some data
    assert_none(&ctx, &storage, k2);
    must_put(&ctx, &storage, k2, v2);

    // isolate leader
    cluster.add_send_filter(IsolationFilterFactory::new(leader.get_store_id()));

    // leader still in lease, check if can read on leader
    assert_eq!(can_read(&ctx, &storage, k2, v2), true);
}

#[test]
fn test_batch_snapshot() {
    let count = 3;
    let mut cluster = new_server_cluster(0, count);
    // Set election timeout and max leader lease to 3s.
    // Large max leader lease stabilizes the test.
    let election_timeout = configure_for_lease_read(&mut cluster, Some(300), Some(10));
    cluster.run();

    let key = b"key";
    // make sure leader has been elected.
    assert_eq!(cluster.must_get(key), None);

    let region = cluster.get_region(b"");
    let leader = cluster.leader_of_region(region.get_id()).unwrap();
    let storage = cluster.sim.rl().storages[&leader.get_id()].clone();

    let mut ctx = Context::new();
    ctx.set_region_id(region.get_id());
    ctx.set_region_epoch(region.get_region_epoch().clone());
    ctx.set_peer(leader.clone());

    let size = 3;
    let batch = vec![ctx.clone(); size];
    let snapshots = must_batch_snapshot(batch, &storage);
    assert_eq!(size, snapshots.len());
    for s in snapshots {
        assert!(s.is_some());
    }
    // sleep util leader lease is expired.
    thread::sleep(election_timeout);
    let batch = vec![ctx; size];
    let snapshots = must_batch_snapshot(batch, &storage);
    assert_eq!(size, snapshots.len());
    for s in snapshots {
        assert!(s.is_none());
    }
}

fn must_batch_snapshot<E: Engine>(batch: Vec<Context>, engine: &E) -> BatchResults<E::Snap> {
    let (tx, rx) = channel();
    let on_finished = box move |snapshots| {
        tx.send(snapshots).unwrap();
    };
    engine.async_batch_snapshot(batch, on_finished).unwrap();
    rx.recv().unwrap()
}

fn must_put<E: Engine>(ctx: &Context, engine: &E, key: &[u8], value: &[u8]) {
    engine.put(ctx, Key::from_raw(key), value.to_vec()).unwrap();
}

fn must_put_cf<E: Engine>(ctx: &Context, engine: &E, cf: CfName, key: &[u8], value: &[u8]) {
    engine
        .put_cf(ctx, cf, Key::from_raw(key), value.to_vec())
        .unwrap();
}

fn must_delete<E: Engine>(ctx: &Context, engine: &E, key: &[u8]) {
    engine.delete(ctx, Key::from_raw(key)).unwrap();
}

fn must_delete_cf<E: Engine>(ctx: &Context, engine: &E, cf: CfName, key: &[u8]) {
    engine.delete_cf(ctx, cf, Key::from_raw(key)).unwrap();
}

fn assert_has<E: Engine>(ctx: &Context, engine: &E, key: &[u8], value: &[u8]) {
    let snapshot = engine.snapshot(ctx).unwrap();
    assert_eq!(snapshot.get(&Key::from_raw(key)).unwrap().unwrap(), value);
}

fn can_read<E: Engine>(ctx: &Context, engine: &E, key: &[u8], value: &[u8]) -> bool {
    if let Ok(s) = engine.snapshot(ctx) {
        assert_eq!(s.get(&Key::from_raw(key)).unwrap().unwrap(), value);
        return true;
    }
    false
}

fn assert_has_cf<E: Engine>(ctx: &Context, engine: &E, cf: CfName, key: &[u8], value: &[u8]) {
    let snapshot = engine.snapshot(ctx).unwrap();
    assert_eq!(
        snapshot.get_cf(cf, &Key::from_raw(key)).unwrap().unwrap(),
        value
    );
}

fn assert_none<E: Engine>(ctx: &Context, engine: &E, key: &[u8]) {
    let snapshot = engine.snapshot(ctx).unwrap();
    assert_eq!(snapshot.get(&Key::from_raw(key)).unwrap(), None);
}

fn assert_none_cf<E: Engine>(ctx: &Context, engine: &E, cf: CfName, key: &[u8]) {
    let snapshot = engine.snapshot(ctx).unwrap();
    assert_eq!(snapshot.get_cf(cf, &Key::from_raw(key)).unwrap(), None);
}

fn assert_seek<E: Engine>(ctx: &Context, engine: &E, key: &[u8], pair: (&[u8], &[u8])) {
    let snapshot = engine.snapshot(ctx).unwrap();
    let mut cursor = snapshot
        .iter(IterOption::default(), ScanMode::Mixed)
        .unwrap();
    let mut statistics = CFStatistics::default();
    cursor.seek(&Key::from_raw(key), &mut statistics).unwrap();
    assert_eq!(cursor.key(&mut statistics), &*bytes::encode_bytes(pair.0));
    assert_eq!(cursor.value(&mut statistics), pair.1);
}

fn assert_seek_cf<E: Engine>(
    ctx: &Context,
    engine: &E,
    cf: CfName,
    key: &[u8],
    pair: (&[u8], &[u8]),
) {
    let snapshot = engine.snapshot(ctx).unwrap();
    let mut cursor = snapshot
        .iter_cf(cf, IterOption::default(), ScanMode::Mixed)
        .unwrap();
    let mut statistics = CFStatistics::default();
    cursor.seek(&Key::from_raw(key), &mut statistics).unwrap();
    assert_eq!(cursor.key(&mut statistics), &*bytes::encode_bytes(pair.0));
    assert_eq!(cursor.value(&mut statistics), pair.1);
}

fn assert_near_seek<I: Iterator>(cursor: &mut Cursor<I>, key: &[u8], pair: (&[u8], &[u8])) {
    let mut statistics = CFStatistics::default();
    assert!(
        cursor
            .near_seek(&Key::from_raw(key), &mut statistics)
            .unwrap(),
        escape(key)
    );
    assert_eq!(cursor.key(&mut statistics), &*bytes::encode_bytes(pair.0));
    assert_eq!(cursor.value(&mut statistics), pair.1);
}

fn assert_near_reverse_seek<I: Iterator>(cursor: &mut Cursor<I>, key: &[u8], pair: (&[u8], &[u8])) {
    let mut statistics = CFStatistics::default();
    assert!(
        cursor
            .near_reverse_seek(&Key::from_raw(key), &mut statistics)
            .unwrap(),
        escape(key)
    );
    assert_eq!(cursor.key(&mut statistics), &*bytes::encode_bytes(pair.0));
    assert_eq!(cursor.value(&mut statistics), pair.1);
}

fn get_put<E: Engine>(ctx: &Context, engine: &E) {
    assert_none(ctx, engine, b"x");
    must_put(ctx, engine, b"x", b"1");
    assert_has(ctx, engine, b"x", b"1");
    must_put(ctx, engine, b"x", b"2");
    assert_has(ctx, engine, b"x", b"2");
}

fn batch<E: Engine>(ctx: &Context, engine: &E) {
    engine
        .write(
            ctx,
            vec![
                Modify::Put(CF_DEFAULT, Key::from_raw(b"x"), b"1".to_vec()),
                Modify::Put(CF_DEFAULT, Key::from_raw(b"y"), b"2".to_vec()),
            ],
        )
        .unwrap();
    assert_has(ctx, engine, b"x", b"1");
    assert_has(ctx, engine, b"y", b"2");

    engine
        .write(
            ctx,
            vec![
                Modify::Delete(CF_DEFAULT, Key::from_raw(b"x")),
                Modify::Delete(CF_DEFAULT, Key::from_raw(b"y")),
            ],
        )
        .unwrap();
    assert_none(ctx, engine, b"y");
    assert_none(ctx, engine, b"y");
}

fn seek<E: Engine>(ctx: &Context, engine: &E) {
    must_put(ctx, engine, b"x", b"1");
    assert_seek(ctx, engine, b"x", (b"x", b"1"));
    assert_seek(ctx, engine, b"a", (b"x", b"1"));
    must_put(ctx, engine, b"z", b"2");
    assert_seek(ctx, engine, b"y", (b"z", b"2"));
    assert_seek(ctx, engine, b"x\x00", (b"z", b"2"));
    let snapshot = engine.snapshot(ctx).unwrap();
    let mut iter = snapshot
        .iter(IterOption::default(), ScanMode::Mixed)
        .unwrap();
    let mut statistics = CFStatistics::default();
    assert!(
        !iter
            .seek(&Key::from_raw(b"z\x00"), &mut statistics)
            .unwrap()
    );
    must_delete(ctx, engine, b"x");
    must_delete(ctx, engine, b"z");
}

fn near_seek<E: Engine>(ctx: &Context, engine: &E) {
    must_put(ctx, engine, b"x", b"1");
    must_put(ctx, engine, b"z", b"2");
    let snapshot = engine.snapshot(ctx).unwrap();
    let mut cursor = snapshot
        .iter(IterOption::default(), ScanMode::Mixed)
        .unwrap();
    assert_near_seek(&mut cursor, b"x", (b"x", b"1"));
    assert_near_seek(&mut cursor, b"a", (b"x", b"1"));
    assert_near_reverse_seek(&mut cursor, b"z1", (b"z", b"2"));
    assert_near_reverse_seek(&mut cursor, b"x1", (b"x", b"1"));
    assert_near_seek(&mut cursor, b"y", (b"z", b"2"));
    assert_near_seek(&mut cursor, b"x\x00", (b"z", b"2"));
    let mut statistics = CFStatistics::default();
    assert!(
        !cursor
            .near_seek(&Key::from_raw(b"z\x00"), &mut statistics)
            .unwrap()
    );
    must_delete(ctx, engine, b"x");
    must_delete(ctx, engine, b"z");
}

fn cf<E: Engine>(ctx: &Context, engine: &E) {
    assert_none_cf(ctx, engine, "default", b"key");
    must_put_cf(ctx, engine, "default", b"key", b"value");
    assert_has_cf(ctx, engine, "default", b"key", b"value");
    assert_seek_cf(ctx, engine, "default", b"k", (b"key", b"value"));
    must_delete_cf(ctx, engine, "default", b"key");
    assert_none_cf(ctx, engine, "default", b"key");
}

fn empty_write<E: Engine>(ctx: &Context, engine: &E) {
    engine.write(ctx, vec![]).unwrap_err();
}

fn wrong_context<E: Engine>(ctx: &Context, engine: &E) {
    let region_id = ctx.get_region_id();
    let mut ctx = ctx.to_owned();
    ctx.set_region_id(region_id + 1);
    assert!(engine.write(&ctx, vec![]).is_err());
}

fn empty_batch_snapshot<E: Engine>(engine: &E) {
    let on_finished = box move |_| {};
    engine
        .async_batch_snapshot(vec![], on_finished)
        .unwrap_err();
}
<|endoftext|>


<commit_before>use super::schema::*;
use diesel::connection::SimpleConnection;
use diesel::*;
use schema_dsl::*;

#[test]
fn selecting_basic_data() {
    use schema::users::dsl::*;

    let connection = connection();
    connection
        .execute("INSERT INTO users (name) VALUES ('Sean'), ('Tess')")
        .unwrap();

    let expected_data = vec![
        ("Sean".to_string(), None::<String>),
        ("Tess".to_string(), None::<String>),
    ];
    let actual_data: Vec<_> = users.select((name, hair_color)).load(&connection).unwrap();
    assert_eq!(expected_data, actual_data);
}

#[test]
fn selecting_a_struct() {
    use schema::users::dsl::*;
    let connection = connection();
    connection
        .execute("INSERT INTO users (name) VALUES ('Sean'), ('Tess')")
        .unwrap();

    let expected_users = vec![NewUser::new("Sean", None), NewUser::new("Tess", None)];
    let actual_users: Vec<_> = users.select((name, hair_color)).load(&connection).unwrap();
    assert_eq!(expected_users, actual_users);
}

#[test]
fn with_safe_select() {
    use schema::users::dsl::*;

    let connection = connection();
    connection
        .execute("INSERT INTO users (name) VALUES ('Sean'), ('Tess')")
        .unwrap();

    let select_name = users.select(name);
    let names: Vec<String> = select_name.load(&connection).unwrap();

    assert_eq!(vec!["Sean".to_string(), "Tess".to_string()], names);
}

#[test]
fn with_select_sql() {
    use diesel::dsl::sql;

    let connection = connection();
    connection
        .execute("INSERT INTO users (name) VALUES ('Sean'), ('Tess')")
        .unwrap();

    let select_count = users::table.select(sql::<sql_types::BigInt>("COUNT(*)"));
    let get_count = || select_count.clone().first::<i64>(&connection);

    assert_eq!(Ok(2), get_count());

    connection
        .execute("INSERT INTO users (name) VALUES ('Jim')")
        .unwrap();

    assert_eq!(Ok(3), get_count());
}

#[test]
fn selecting_nullable_followed_by_non_null() {
    use schema::users::dsl::*;

    let connection = connection();
    connection
        .execute("INSERT INTO users (name) VALUES ('Sean')")
        .unwrap();

    let source = users.select((hair_color, name));
    let expected_data = vec![(None::<String>, "Sean".to_string())];
    let data: Vec<_> = source.load(&connection).unwrap();

    assert_eq!(expected_data, data);
}

#[test]
fn selecting_expression_with_bind_param() {
    use schema::users::dsl::*;

    let connection = connection();
    connection
        .execute("INSERT INTO users (name) VALUES ('Sean'), ('Tess')")
        .unwrap();

    let source = users.select(name.eq("Sean".to_string()));
    let expected_data = vec![true, false];
    let actual_data = source.load::<bool>(&connection).unwrap();

    assert_eq!(expected_data, actual_data);
}

table! {
    select {
        id -> Integer,
        join -> Integer,
    }
}

#[test]
#[cfg(not(feature = "mysql"))] // FIXME: Figure out how to handle tests that modify schema
fn selecting_columns_and_tables_with_reserved_names() {
    use self::select::dsl::*;

    let connection = connection();
    create_table(
        "select",
        (
            integer("id").primary_key().auto_increment(),
            integer("join").not_null(),
        ),
    ).execute(&connection)
    .unwrap();
    connection
        .execute("INSERT INTO \"select\" (\"join\") VALUES (1), (2), (3)")
        .unwrap();

    let expected_data = vec![(1, 1), (2, 2), (3, 3)];
    let actual_data: Vec<(i32, i32)> = select.load(&connection).unwrap();
    assert_eq!(expected_data, actual_data);

    let expected_data = vec![1, 2, 3];
    let actual_data: Vec<i32> = select.select(join).load(&connection).unwrap();
    assert_eq!(expected_data, actual_data);
}

#[test]
#[cfg(not(feature = "mysql"))] // FIXME: Figure out how to handle tests that modify schema
fn selecting_columns_with_different_definition_order() {
    let connection = connection();
    drop_table_cascade(&connection, "users");
    create_table(
        "users",
        (
            integer("id").primary_key().auto_increment(),
            string("hair_color"),
            string("name").not_null(),
        ),
    ).execute(&connection)
    .unwrap();
    let expected_user = User::with_hair_color(1, "Sean", "black");
    insert_into(users::table)
        .values(&NewUser::new("Sean", Some("black")))
        .execute(&connection)
        .unwrap();
    let user_from_select = users::table.first(&connection);

    assert_eq!(Ok(&expected_user), user_from_select.as_ref());
}

#[test]
fn selection_using_subselect() {
    use schema::posts::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let ids: Vec<i32> = users::table.select(users::id).load(&connection).unwrap();
    let query = format!(
        "INSERT INTO posts (user_id, title) VALUES ({}, 'Hello'), ({}, 'World')",
        ids[0], ids[1]
    );
    connection.execute(&query).unwrap();

    let users = users::table
        .filter(users::name.eq("Sean"))
        .select(users::id);
    let data: Vec<String> = posts
        .select(title)
        .filter(user_id.eq_any(users))
        .load(&connection)
        .unwrap();

    assert_eq!(vec!["Hello".to_string()], data);
}

table! {
    users_select_for_update_modifieres {
        id -> Integer,
        name -> Text,
        hair_color -> Nullable<Text>,
    }
}

table! {
    users_select_for_update {
        id -> Integer,
        name -> Text,
        hair_color -> Nullable<Text>,
    }
}

table! {
    users_select_for_no_key_update {
        id -> Integer,
        name -> Text,
        hair_color -> Nullable<Text>,
    }
}

table! {
    users_fk_for_no_key_update {
        id -> Integer,
        users_fk -> Integer,
    }
}

#[cfg(not(feature = "sqlite"))]
#[test]
fn select_for_update_locks_selected_rows() {
    use self::users_select_for_update::dsl::*;
    use std::mem::drop;
    use std::sync::mpsc;
    use std::thread;
    use std::time::Duration;

    let conn_1 = connection_without_transaction();
    conn_1
        .execute("DROP TABLE IF EXISTS users_select_for_update")
        .unwrap();
    create_table(
        "users_select_for_update",
        (
            integer("id").primary_key().auto_increment(),
            string("name").not_null(),
            string("hair_color"),
        ),
    ).execute(&conn_1)
    .unwrap();
    conn_1
        .batch_execute(
            "
            -- MySQL locks the whole table without this index
            CREATE UNIQUE INDEX users_select_for_update_name ON users_select_for_update (name);
            INSERT INTO users_select_for_update (name) VALUES ('Sean'), ('Tess');
        ",
        ).unwrap();
    conn_1.begin_test_transaction().unwrap();

    let _sean = users_select_for_update
        .for_update()
        .filter(name.eq("Sean"))
        .first::<User>(&conn_1)
        .unwrap();

    let (send, recv) = mpsc::channel();
    let send2 = send.clone();

    let _blocked_thread = thread::spawn(move || {
        let conn_2 = connection();
        update(users_select_for_update.filter(name.eq("Sean")))
            .set(name.eq("Jim"))
            .execute(&conn_2)
            .unwrap();
        send.send("Sean").unwrap();
    });

    let _unblocked_thread = thread::spawn(move || {
        let conn_3 = connection();
        update(users_select_for_update.filter(name.eq("Tess")))
            .set(name.eq("Bob"))
            .execute(&conn_3)
            .unwrap();
        send2.send("Tess").unwrap();
    });

    let timeout = Duration::from_secs(1);
    let next_selected_name = recv.recv_timeout(timeout).unwrap();
    // conn_3 will always complete before conn_2, as conn_2 is blocked
    assert_eq!("Tess", next_selected_name);
    let next_selected_name = recv.recv_timeout(timeout);
    // conn_2 should still be blocked
    assert!(next_selected_name.is_err());
    drop(conn_1);
    let next_selected_name = recv.recv_timeout(timeout).unwrap();
    // Dropping conn_1 unblocks conn_2
    assert_eq!("Sean", next_selected_name);
}

#[cfg(feature = "postgres")]
#[test]
fn select_for_update_modifiers() {
    use self::users_select_for_update_modifieres::dsl::*;

    // We need to actually commit some data for the
    // test
    let conn_1 = connection_without_transaction();
    let conn_2 = connection();
    let conn_3 = connection();

    // Recreate the table
    conn_1
        .execute("DROP TABLE IF EXISTS users_select_for_update_modifieres")
        .unwrap();
    create_table(
        "users_select_for_update_modifieres",
        (
            integer("id").primary_key().auto_increment(),
            string("name").not_null(),
            string("hair_color"),
        ),
    ).execute(&conn_1)
    .unwrap();

    // Add some test data
    conn_1
        .execute(
            "
            INSERT INTO users_select_for_update_modifieres (name)
            VALUES ('Sean'), ('Tess')
            ",
        )
        .unwrap();

    // Now both connections have begun a transaction
    conn_1.begin_test_transaction().unwrap();

    // Lock the "Sean" row
    let _sean = users_select_for_update_modifieres
        .order(name)
        .for_update()
        .first::<User>(&conn_1)
        .unwrap();

    // Try to access the "Sean" row with `NOWAIT`
    conn_2.execute("SET STATEMENT_TIMEOUT TO 1000").unwrap();
    let result = users_select_for_update_modifieres
        .order(name)
        .for_update()
        .no_wait()
        .first::<User>(&conn_2);

    // Make sure we errored in the correct way (without timing out)
    assert!(result.is_err());
    if !format!("{:?}", result).contains("could not obtain lock on row") {
        panic!("{:?}", result);
    }

    // Try to access the "Sean" row with `SKIP LOCKED`
    let tess = users_select_for_update_modifieres
        .order(name)
        .for_update()
        .skip_locked()
        .first::<User>(&conn_3)
        .unwrap();

    // Make sure got back "Tess"
    assert_eq!(tess.name, "Tess");
}

#[cfg(feature = "postgres")]
#[test]
fn select_for_no_key_update_modifiers() {
    use self::users_fk_for_no_key_update::dsl::*;
    use self::users_select_for_no_key_update::dsl::*;

    // We need to actually commit some data for the
    // test
    let conn_1 = connection_without_transaction();
    let conn_2 = connection();
    let conn_3 = connection();
    let conn_4 = connection();

    // Recreate the table
    conn_1
        .execute("DROP TABLE IF EXISTS users_fk_for_no_key_update")
        .unwrap();
    conn_1
        .execute("DROP TABLE IF EXISTS users_select_for_no_key_update")
        .unwrap();

    create_table(
        "users_select_for_no_key_update",
        (
            integer("id").primary_key().auto_increment(),
            string("name").not_null(),
            string("hair_color"),
        ),
    ).execute(&conn_1)
    .unwrap();

    create_table(
        "users_fk_for_no_key_update",
        (
            integer("id").primary_key().auto_increment(),
            integer("users_fk").not_null(),
        ),
    ).execute(&conn_1)
    .unwrap();

    // Add a foreign key
    conn_1
        .execute(
            "ALTER TABLE users_fk_for_no_key_update ADD CONSTRAINT users_fk \
             FOREIGN KEY (users_fk) REFERENCES users_select_for_no_key_update(id)",
        ).unwrap();

    // Add some test data
    conn_1
        .execute(
            "INSERT INTO users_select_for_no_key_update (name) VALUES ('Sean'), ('Tess'), ('Will')",
        ).unwrap();

    conn_1.begin_test_transaction().unwrap();

    // Lock the "Sean" row, except the key
    let _sean = users_select_for_no_key_update
        .order(name)
        .for_no_key_update()
        .first::<User>(&conn_1)
        .unwrap();

    // Try to add an object referencing the "Sean" row
    conn_2
        .execute(
            "INSERT INTO users_fk_for_no_key_update (users_fk) \
             SELECT id FROM users_select_for_no_key_update where name='Sean'",
        ).unwrap();

    // Check that it was successfully added
    let expected_data = vec![(1, 1)];
    let actual_data: Vec<(i32, i32)> = users_fk_for_no_key_update.load(&conn_2).unwrap();
    assert_eq!(expected_data, actual_data);

    // Try to access the "Sean" row with `for no key update` and `SKIP LOCKED`
    let tess = users_select_for_no_key_update
        .order(name)
        .for_no_key_update()
        .skip_locked()
        .first::<User>(&conn_3)
        .unwrap();

    // Make sure got back "Tess"
    assert_eq!(tess.name, "Tess");

    // Lock the "Will" row completely
    let will = users_select_for_no_key_update
        .order(name)
        .for_update()
        .skip_locked()
        .first::<User>(&conn_4)
        .unwrap();

    assert_eq!(will.name, "Will");

    conn_2.execute("SET STATEMENT_TIMEOUT TO 1000").unwrap();
    let result = conn_2.execute(
        "INSERT INTO users_fk_for_no_key_update (users_fk) \
         SELECT id FROM users_select_for_no_key_update where name='Will'",
    );

    // Times out instead of inserting row
    assert!(result.is_err());
    if !format!("{:?}", result).contains("canceling statement due to statement timeout") {
        panic!("{:?}", result);
    }
}

#[test]
fn select_can_be_called_on_query_that_is_valid_subselect_but_invalid_query() {
    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);
    let tess = find_user_by_name("Tess", &connection);
    insert_into(posts::table)
        .values(&vec![
            tess.new_post("Tess", None),
            sean.new_post("Hi", None),
        ]).execute(&connection)
        .unwrap();

    let invalid_query_but_valid_subselect = posts::table
        .filter(posts::title.eq(users::name))
        .select(posts::user_id);
    let users_with_post_using_name_as_title = users::table
        .filter(users::id.eq_any(invalid_query_but_valid_subselect))
        .load(&connection);

    assert_eq!(Ok(vec![tess]), users_with_post_using_name_as_title);
}
<commit_msg>Fix rustfmt<commit_after>use super::schema::*;
use diesel::connection::SimpleConnection;
use diesel::*;
use schema_dsl::*;

#[test]
fn selecting_basic_data() {
    use schema::users::dsl::*;

    let connection = connection();
    connection
        .execute("INSERT INTO users (name) VALUES ('Sean'), ('Tess')")
        .unwrap();

    let expected_data = vec![
        ("Sean".to_string(), None::<String>),
        ("Tess".to_string(), None::<String>),
    ];
    let actual_data: Vec<_> = users.select((name, hair_color)).load(&connection).unwrap();
    assert_eq!(expected_data, actual_data);
}

#[test]
fn selecting_a_struct() {
    use schema::users::dsl::*;
    let connection = connection();
    connection
        .execute("INSERT INTO users (name) VALUES ('Sean'), ('Tess')")
        .unwrap();

    let expected_users = vec![NewUser::new("Sean", None), NewUser::new("Tess", None)];
    let actual_users: Vec<_> = users.select((name, hair_color)).load(&connection).unwrap();
    assert_eq!(expected_users, actual_users);
}

#[test]
fn with_safe_select() {
    use schema::users::dsl::*;

    let connection = connection();
    connection
        .execute("INSERT INTO users (name) VALUES ('Sean'), ('Tess')")
        .unwrap();

    let select_name = users.select(name);
    let names: Vec<String> = select_name.load(&connection).unwrap();

    assert_eq!(vec!["Sean".to_string(), "Tess".to_string()], names);
}

#[test]
fn with_select_sql() {
    use diesel::dsl::sql;

    let connection = connection();
    connection
        .execute("INSERT INTO users (name) VALUES ('Sean'), ('Tess')")
        .unwrap();

    let select_count = users::table.select(sql::<sql_types::BigInt>("COUNT(*)"));
    let get_count = || select_count.clone().first::<i64>(&connection);

    assert_eq!(Ok(2), get_count());

    connection
        .execute("INSERT INTO users (name) VALUES ('Jim')")
        .unwrap();

    assert_eq!(Ok(3), get_count());
}

#[test]
fn selecting_nullable_followed_by_non_null() {
    use schema::users::dsl::*;

    let connection = connection();
    connection
        .execute("INSERT INTO users (name) VALUES ('Sean')")
        .unwrap();

    let source = users.select((hair_color, name));
    let expected_data = vec![(None::<String>, "Sean".to_string())];
    let data: Vec<_> = source.load(&connection).unwrap();

    assert_eq!(expected_data, data);
}

#[test]
fn selecting_expression_with_bind_param() {
    use schema::users::dsl::*;

    let connection = connection();
    connection
        .execute("INSERT INTO users (name) VALUES ('Sean'), ('Tess')")
        .unwrap();

    let source = users.select(name.eq("Sean".to_string()));
    let expected_data = vec![true, false];
    let actual_data = source.load::<bool>(&connection).unwrap();

    assert_eq!(expected_data, actual_data);
}

table! {
    select {
        id -> Integer,
        join -> Integer,
    }
}

#[test]
#[cfg(not(feature = "mysql"))] // FIXME: Figure out how to handle tests that modify schema
fn selecting_columns_and_tables_with_reserved_names() {
    use self::select::dsl::*;

    let connection = connection();
    create_table(
        "select",
        (
            integer("id").primary_key().auto_increment(),
            integer("join").not_null(),
        ),
    ).execute(&connection)
    .unwrap();
    connection
        .execute("INSERT INTO \"select\" (\"join\") VALUES (1), (2), (3)")
        .unwrap();

    let expected_data = vec![(1, 1), (2, 2), (3, 3)];
    let actual_data: Vec<(i32, i32)> = select.load(&connection).unwrap();
    assert_eq!(expected_data, actual_data);

    let expected_data = vec![1, 2, 3];
    let actual_data: Vec<i32> = select.select(join).load(&connection).unwrap();
    assert_eq!(expected_data, actual_data);
}

#[test]
#[cfg(not(feature = "mysql"))] // FIXME: Figure out how to handle tests that modify schema
fn selecting_columns_with_different_definition_order() {
    let connection = connection();
    drop_table_cascade(&connection, "users");
    create_table(
        "users",
        (
            integer("id").primary_key().auto_increment(),
            string("hair_color"),
            string("name").not_null(),
        ),
    ).execute(&connection)
    .unwrap();
    let expected_user = User::with_hair_color(1, "Sean", "black");
    insert_into(users::table)
        .values(&NewUser::new("Sean", Some("black")))
        .execute(&connection)
        .unwrap();
    let user_from_select = users::table.first(&connection);

    assert_eq!(Ok(&expected_user), user_from_select.as_ref());
}

#[test]
fn selection_using_subselect() {
    use schema::posts::dsl::*;

    let connection = connection_with_sean_and_tess_in_users_table();
    let ids: Vec<i32> = users::table.select(users::id).load(&connection).unwrap();
    let query = format!(
        "INSERT INTO posts (user_id, title) VALUES ({}, 'Hello'), ({}, 'World')",
        ids[0], ids[1]
    );
    connection.execute(&query).unwrap();

    let users = users::table
        .filter(users::name.eq("Sean"))
        .select(users::id);
    let data: Vec<String> = posts
        .select(title)
        .filter(user_id.eq_any(users))
        .load(&connection)
        .unwrap();

    assert_eq!(vec!["Hello".to_string()], data);
}

table! {
    users_select_for_update_modifieres {
        id -> Integer,
        name -> Text,
        hair_color -> Nullable<Text>,
    }
}

table! {
    users_select_for_update {
        id -> Integer,
        name -> Text,
        hair_color -> Nullable<Text>,
    }
}

table! {
    users_select_for_no_key_update {
        id -> Integer,
        name -> Text,
        hair_color -> Nullable<Text>,
    }
}

table! {
    users_fk_for_no_key_update {
        id -> Integer,
        users_fk -> Integer,
    }
}

#[cfg(not(feature = "sqlite"))]
#[test]
fn select_for_update_locks_selected_rows() {
    use self::users_select_for_update::dsl::*;
    use std::mem::drop;
    use std::sync::mpsc;
    use std::thread;
    use std::time::Duration;

    let conn_1 = connection_without_transaction();
    conn_1
        .execute("DROP TABLE IF EXISTS users_select_for_update")
        .unwrap();
    create_table(
        "users_select_for_update",
        (
            integer("id").primary_key().auto_increment(),
            string("name").not_null(),
            string("hair_color"),
        ),
    ).execute(&conn_1)
    .unwrap();
    conn_1
        .batch_execute(
            "
            -- MySQL locks the whole table without this index
            CREATE UNIQUE INDEX users_select_for_update_name ON users_select_for_update (name);
            INSERT INTO users_select_for_update (name) VALUES ('Sean'), ('Tess');
        ",
        ).unwrap();
    conn_1.begin_test_transaction().unwrap();

    let _sean = users_select_for_update
        .for_update()
        .filter(name.eq("Sean"))
        .first::<User>(&conn_1)
        .unwrap();

    let (send, recv) = mpsc::channel();
    let send2 = send.clone();

    let _blocked_thread = thread::spawn(move || {
        let conn_2 = connection();
        update(users_select_for_update.filter(name.eq("Sean")))
            .set(name.eq("Jim"))
            .execute(&conn_2)
            .unwrap();
        send.send("Sean").unwrap();
    });

    let _unblocked_thread = thread::spawn(move || {
        let conn_3 = connection();
        update(users_select_for_update.filter(name.eq("Tess")))
            .set(name.eq("Bob"))
            .execute(&conn_3)
            .unwrap();
        send2.send("Tess").unwrap();
    });

    let timeout = Duration::from_secs(1);
    let next_selected_name = recv.recv_timeout(timeout).unwrap();
    // conn_3 will always complete before conn_2, as conn_2 is blocked
    assert_eq!("Tess", next_selected_name);
    let next_selected_name = recv.recv_timeout(timeout);
    // conn_2 should still be blocked
    assert!(next_selected_name.is_err());
    drop(conn_1);
    let next_selected_name = recv.recv_timeout(timeout).unwrap();
    // Dropping conn_1 unblocks conn_2
    assert_eq!("Sean", next_selected_name);
}

#[cfg(feature = "postgres")]
#[test]
fn select_for_update_modifiers() {
    use self::users_select_for_update_modifieres::dsl::*;

    // We need to actually commit some data for the
    // test
    let conn_1 = connection_without_transaction();
    let conn_2 = connection();
    let conn_3 = connection();

    // Recreate the table
    conn_1
        .execute("DROP TABLE IF EXISTS users_select_for_update_modifieres")
        .unwrap();
    create_table(
        "users_select_for_update_modifieres",
        (
            integer("id").primary_key().auto_increment(),
            string("name").not_null(),
            string("hair_color"),
        ),
    ).execute(&conn_1)
    .unwrap();

    // Add some test data
    conn_1
        .execute(
            "
            INSERT INTO users_select_for_update_modifieres (name)
            VALUES ('Sean'), ('Tess')
            ",
        ).unwrap();

    // Now both connections have begun a transaction
    conn_1.begin_test_transaction().unwrap();

    // Lock the "Sean" row
    let _sean = users_select_for_update_modifieres
        .order(name)
        .for_update()
        .first::<User>(&conn_1)
        .unwrap();

    // Try to access the "Sean" row with `NOWAIT`
    conn_2.execute("SET STATEMENT_TIMEOUT TO 1000").unwrap();
    let result = users_select_for_update_modifieres
        .order(name)
        .for_update()
        .no_wait()
        .first::<User>(&conn_2);

    // Make sure we errored in the correct way (without timing out)
    assert!(result.is_err());
    if !format!("{:?}", result).contains("could not obtain lock on row") {
        panic!("{:?}", result);
    }

    // Try to access the "Sean" row with `SKIP LOCKED`
    let tess = users_select_for_update_modifieres
        .order(name)
        .for_update()
        .skip_locked()
        .first::<User>(&conn_3)
        .unwrap();

    // Make sure got back "Tess"
    assert_eq!(tess.name, "Tess");
}

#[cfg(feature = "postgres")]
#[test]
fn select_for_no_key_update_modifiers() {
    use self::users_fk_for_no_key_update::dsl::*;
    use self::users_select_for_no_key_update::dsl::*;

    // We need to actually commit some data for the
    // test
    let conn_1 = connection_without_transaction();
    let conn_2 = connection();
    let conn_3 = connection();
    let conn_4 = connection();

    // Recreate the table
    conn_1
        .execute("DROP TABLE IF EXISTS users_fk_for_no_key_update")
        .unwrap();
    conn_1
        .execute("DROP TABLE IF EXISTS users_select_for_no_key_update")
        .unwrap();

    create_table(
        "users_select_for_no_key_update",
        (
            integer("id").primary_key().auto_increment(),
            string("name").not_null(),
            string("hair_color"),
        ),
    ).execute(&conn_1)
    .unwrap();

    create_table(
        "users_fk_for_no_key_update",
        (
            integer("id").primary_key().auto_increment(),
            integer("users_fk").not_null(),
        ),
    ).execute(&conn_1)
    .unwrap();

    // Add a foreign key
    conn_1
        .execute(
            "ALTER TABLE users_fk_for_no_key_update ADD CONSTRAINT users_fk \
             FOREIGN KEY (users_fk) REFERENCES users_select_for_no_key_update(id)",
        ).unwrap();

    // Add some test data
    conn_1
        .execute(
            "INSERT INTO users_select_for_no_key_update (name) VALUES ('Sean'), ('Tess'), ('Will')",
        ).unwrap();

    conn_1.begin_test_transaction().unwrap();

    // Lock the "Sean" row, except the key
    let _sean = users_select_for_no_key_update
        .order(name)
        .for_no_key_update()
        .first::<User>(&conn_1)
        .unwrap();

    // Try to add an object referencing the "Sean" row
    conn_2
        .execute(
            "INSERT INTO users_fk_for_no_key_update (users_fk) \
             SELECT id FROM users_select_for_no_key_update where name='Sean'",
        ).unwrap();

    // Check that it was successfully added
    let expected_data = vec![(1, 1)];
    let actual_data: Vec<(i32, i32)> = users_fk_for_no_key_update.load(&conn_2).unwrap();
    assert_eq!(expected_data, actual_data);

    // Try to access the "Sean" row with `for no key update` and `SKIP LOCKED`
    let tess = users_select_for_no_key_update
        .order(name)
        .for_no_key_update()
        .skip_locked()
        .first::<User>(&conn_3)
        .unwrap();

    // Make sure got back "Tess"
    assert_eq!(tess.name, "Tess");

    // Lock the "Will" row completely
    let will = users_select_for_no_key_update
        .order(name)
        .for_update()
        .skip_locked()
        .first::<User>(&conn_4)
        .unwrap();

    assert_eq!(will.name, "Will");

    conn_2.execute("SET STATEMENT_TIMEOUT TO 1000").unwrap();
    let result = conn_2.execute(
        "INSERT INTO users_fk_for_no_key_update (users_fk) \
         SELECT id FROM users_select_for_no_key_update where name='Will'",
    );

    // Times out instead of inserting row
    assert!(result.is_err());
    if !format!("{:?}", result).contains("canceling statement due to statement timeout") {
        panic!("{:?}", result);
    }
}

#[test]
fn select_can_be_called_on_query_that_is_valid_subselect_but_invalid_query() {
    let connection = connection_with_sean_and_tess_in_users_table();
    let sean = find_user_by_name("Sean", &connection);
    let tess = find_user_by_name("Tess", &connection);
    insert_into(posts::table)
        .values(&vec![
            tess.new_post("Tess", None),
            sean.new_post("Hi", None),
        ]).execute(&connection)
        .unwrap();

    let invalid_query_but_valid_subselect = posts::table
        .filter(posts::title.eq(users::name))
        .select(posts::user_id);
    let users_with_post_using_name_as_title = users::table
        .filter(users::id.eq_any(invalid_query_but_valid_subselect))
        .load(&connection);

    assert_eq!(Ok(vec![tess]), users_with_post_using_name_as_title);
}
<|endoftext|>


<commit_before>pub fn get_array() {}

#[cfg(test)]
mod tests {
    #[test]
    fn test_get_array() {
        assert!(true)
    }
}
<commit_msg>feat: add workaround method for Config::get_array()<commit_after>use config::{Config, Value};

/// Gets an array from a config file using a key, returning `None` if
/// the key wasn't present or couldn't be coerced into a `Vec<Value>`.
///
/// `Config::get_array()` in version `0.4.1` has a bug where it
/// consumes `self` instead of taking `self` by reference. This
/// function is a workaround for the time being. The bug has been
/// fixed but will not be available until the release of version `0.5`
/// of `config-rs`.
pub fn ref_get_array(key: &str, config: &Config) -> Option<Vec<Value>> {
    config.get(key).and_then(Value::into_array)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    // The real test is that this compiles successfully without giving
    // a warning about the config file being moved.
    fn test_ref_get_array() {
        let mut c: Config = Config::new();
        c.set("numbers", vec![1, 2, 3]).unwrap();
        c.set("boolean", true).unwrap();

        // Uncomment and the compiler should complain
        // assert_eq!(c.get_array("numbers"),
        //            Some(vec![Value::Integer(1), Value::Integer(2), Value::Integer(3)]));
        // assert_eq!(c.get_array("boolean"), None); // Use of moved value: c

        assert_eq!(ref_get_array("numbers", &c),
                   Some(vec![Value::Integer(1), Value::Integer(2), Value::Integer(3)]));
        assert_eq!(ref_get_array("boolean", &c), None);
    }
}
<|endoftext|>