edgedb-protocol 0.6.0

use std::fs;
use std::collections::HashMap;
use std::error::Error;

use uuid::Uuid;
use bytes::{Bytes, BytesMut};

use edgedb_protocol::common::{Capabilities, RawTypedesc};
use edgedb_protocol::encoding::{Input, Output};
use edgedb_protocol::features::ProtocolVersion;
use edgedb_protocol::server_message::{Authentication};
use edgedb_protocol::server_message::{CommandComplete0, CommandComplete1};
use edgedb_protocol::server_message::{CommandDataDescription0, Data};
use edgedb_protocol::server_message::{CommandDataDescription1};
use edgedb_protocol::server_message::{ErrorResponse, ErrorSeverity};
use edgedb_protocol::server_message::{LogMessage, MessageSeverity};
use edgedb_protocol::server_message::{PrepareComplete, Cardinality};
use edgedb_protocol::server_message::{ReadyForCommand, TransactionState};
use edgedb_protocol::server_message::{RestoreReady};
use edgedb_protocol::server_message::{ServerHandshake};
use edgedb_protocol::server_message::{ServerKeyData, ParameterStatus};
use edgedb_protocol::server_message::{ServerMessage};
use edgedb_protocol::server_message::{StateDataDescription};

mod base;

macro_rules! encoding_eq_ver {
    ($major: expr, $minor: expr, $message: expr, $bytes: expr) => {
        let proto = ProtocolVersion::new($major, $minor);
        let data: &[u8] = $bytes;
        let mut bytes = BytesMut::new();
        $message.encode(&mut Output::new(&proto, &mut bytes))?;
        println!("Serialized bytes {:?}", bytes);
        let bytes = bytes.freeze();
        assert_eq!(&bytes[..], data);
        assert_eq!(
            ServerMessage::decode(
                &mut Input::new(proto, Bytes::copy_from_slice(data))
            )?,
            $message,
        );
    }
}

macro_rules! encoding_eq {
    ($message: expr, $bytes: expr) => {
        let (major, minor) = ProtocolVersion::current().version_tuple();
        encoding_eq_ver!(major, minor, $message, $bytes);
    }
}

macro_rules! map {
    ($($key:expr => $value:expr),*) => {
        {
            #[allow(unused_mut)]
            let mut h = HashMap::new();
            $(
                h.insert($key, $value);
            )*
            h
        }
    }
}


#[test]
fn server_handshake() -> Result<(), Box<dyn Error>> {
    encoding_eq!(ServerMessage::ServerHandshake(ServerHandshake {
        major_ver: 1,
        minor_ver: 0,
        extensions: HashMap::new(),
    }), b"v\0\0\0\n\0\x01\0\0\0\0");
    Ok(())
}

#[test]
fn ready_for_command() -> Result<(), Box<dyn Error>> {
    encoding_eq!(ServerMessage::ReadyForCommand(ReadyForCommand {
        transaction_state: TransactionState::NotInTransaction,
        headers: HashMap::new(),
    }), b"Z\0\0\0\x07\0\0I");
    Ok(())
}

#[test]
fn error_response() -> Result<(), Box<dyn Error>> {
    encoding_eq!(ServerMessage::ErrorResponse(ErrorResponse {
        severity: ErrorSeverity::Error,
        code: 50397184,
        message: String::from("missing required connection parameter \
                               in ClientHandshake message: \"user\""),
        attributes: map!{
            257 => Bytes::from_static("Traceback (most recent call last):\n  File \"edb/server/mng_port/edgecon.pyx\", line 1077, in edb.server.mng_port.edgecon.EdgeConnection.main\n    await self.auth()\n  File \"edb/server/mng_port/edgecon.pyx\", line 178, in auth\n    raise errors.BinaryProtocolError(\nedb.errors.BinaryProtocolError: missing required connection parameter in ClientHandshake message: \"user\"\n".as_bytes())
        },
    }), &fs::read("tests/error_response.bin")?[..]);
    Ok(())
}

#[test]
fn server_key_data() -> Result<(), Box<dyn Error>> {
    encoding_eq!(ServerMessage::ServerKeyData(ServerKeyData {
        data: [0u8; 32],
    }), &fs::read("tests/server_key_data.bin")?[..]);
    Ok(())
}

#[test]
fn parameter_status() -> Result<(), Box<dyn Error>> {
    encoding_eq!(ServerMessage::ParameterStatus(ParameterStatus {
        proto: ProtocolVersion::current(),
        name: Bytes::from_static(b"pgaddr"),
        value: Bytes::from_static(b"/work/tmp/db/.s.PGSQL.60128"),
    }), &fs::read("tests/parameter_status.bin")?[..]);
    Ok(())
}

#[test]
fn command_complete0() -> Result<(), Box<dyn Error>> {
    encoding_eq_ver!(0, 13, ServerMessage::CommandComplete0(CommandComplete0 {
        headers: HashMap::new(),
        status_data: Bytes::from_static(b"okay"),
    }), b"C\0\0\0\x0e\0\0\0\0\0\x04okay");
    Ok(())
}

#[test]
fn command_complete1() -> Result<(), Box<dyn Error>> {
    encoding_eq_ver!(1, 0, ServerMessage::CommandComplete1(CommandComplete1 {
        annotations: HashMap::new(),
        capabilities: Capabilities::MODIFICATIONS,
        status_data: Bytes::from_static(b"okay"),
        state: None,
    }), b"C\0\0\0*\0\0\0\0\0\0\0\0\0\x01\0\0\0\x04okay\
          \0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\
          \0\0\0\0");
    Ok(())
}

#[test]
fn prepare_complete() -> Result<(), Box<dyn Error>> {
    encoding_eq!(ServerMessage::PrepareComplete(PrepareComplete {
        headers: HashMap::new(),
        cardinality: Cardinality::AtMostOne,
        input_typedesc_id: Uuid::from_u128(0xFF),
        output_typedesc_id: Uuid::from_u128(0x105),
    }), b"1\0\0\0'\0\0o\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x01\x05");
    encoding_eq!(ServerMessage::PrepareComplete(PrepareComplete {
        headers: HashMap::new(),
        cardinality: Cardinality::NoResult,
        input_typedesc_id: Uuid::from_u128(0xFF),
        output_typedesc_id: Uuid::from_u128(0x0),
    }), b"1\0\0\0'\0\0n\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0");
    Ok(())
}

#[test]
fn command_data_description0() -> Result<(), Box<dyn Error>> {
    encoding_eq_ver!(0, 13,
        ServerMessage::CommandDataDescription0(CommandDataDescription0 {
            headers: HashMap::new(),
            result_cardinality: Cardinality::AtMostOne,
            input: RawTypedesc {
                proto: ProtocolVersion::new(0, 13),
                id: Uuid::from_u128(0xFF),
                data: Bytes::from_static(
                    b"\x04\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff\0\0"),
            },
            output: RawTypedesc {
                proto: ProtocolVersion::new(0, 13),
                id: Uuid::from_u128(0x105),
                data: Bytes::from_static(
                    b"\x02\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x01\x05"),
            },
        }), bconcat!(b"T\0\0\0S\0\0o"
                     b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff"
                     b"\0\0\0\x13"
                     b"\x04\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff\0"
                     b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x01\x05"
                     b"\0\0\0\x11"
                     b"\x02\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x01\x05"));
    encoding_eq_ver!(0, 13,
        ServerMessage::CommandDataDescription0(CommandDataDescription0 {
            headers: HashMap::new(),
            result_cardinality: Cardinality::NoResult,
            input: RawTypedesc {
                proto: ProtocolVersion::new(0, 13),
                id: Uuid::from_u128(0xFF),
                data: Bytes::from_static(
                    b"\x04\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff\0\0"),
            },
            output: RawTypedesc {
                proto: ProtocolVersion::new(0, 13),
                id: Uuid::from_u128(0),
                data: Bytes::from_static(b""),
            },
        }), bconcat!(b"T\0\0\0B\0\0n"
                     b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff"
                     b"\0\0\0\x13"
                     b"\x04\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff\0"
                     b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0"));
    Ok(())
}

#[test]
fn command_data_description1() -> Result<(), Box<dyn Error>> {
    encoding_eq_ver!(1, 0,
        ServerMessage::CommandDataDescription1(CommandDataDescription1 {
            annotations: HashMap::new(),
            capabilities: Capabilities::MODIFICATIONS,
            result_cardinality: Cardinality::AtMostOne,
            input: RawTypedesc {
                proto: ProtocolVersion::current(),
                id: Uuid::from_u128(0xFF),
                data: Bytes::from_static(
                    b"\x04\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff\0\0"),
            },
            output: RawTypedesc {
                proto: ProtocolVersion::current(),
                id: Uuid::from_u128(0x105),
                data: Bytes::from_static(
                    b"\x02\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x01\x05"),
            },
        }), bconcat!(b"T\0\0\0[\0\0\0\0\0\0\0\0\0\x01o"
                     b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff"
                     b"\0\0\0\x13"
                     b"\x04\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff\0"
                     b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x01\x05"
                     b"\0\0\0\x11"
                     b"\x02\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x01\x05"));
    encoding_eq_ver!(1, 0,
        ServerMessage::CommandDataDescription1(CommandDataDescription1 {
            annotations: HashMap::new(),
            capabilities: Capabilities::MODIFICATIONS,
            result_cardinality: Cardinality::NoResult,
            input: RawTypedesc {
                proto: ProtocolVersion::current(),
                id: Uuid::from_u128(0xFF),
                data: Bytes::from_static(
                    b"\x04\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff\0\0"),
            },
            output: RawTypedesc {
                proto: ProtocolVersion::current(),
                id: Uuid::from_u128(0),
                data: Bytes::from_static(b""),
            },
        }), bconcat!(b"T\0\0\0J\0\0\0\0\0\0\0\0\0\x01n"
                     b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff"
                     b"\0\0\0\x13"
                     b"\x04\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\xff\0"
                     b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0"));
    Ok(())
}

#[test]
fn data() -> Result<(), Box<dyn Error>> {
    encoding_eq!(ServerMessage::Data(Data {
        data: vec![Bytes::from_static(b"\0\0\0\0\0\0\0\x01")],
    }), b"D\0\0\0\x12\0\x01\0\0\0\x08\0\0\0\0\0\0\0\x01");
    Ok(())
}

#[test]
fn restore_ready() -> Result<(), Box<dyn Error>> {
    encoding_eq!(ServerMessage::RestoreReady(RestoreReady {
        jobs: 1,
        headers: HashMap::new(),
    }), b"+\0\0\0\x08\0\0\0\x01");
    Ok(())
}

#[test]
fn authentication() -> Result<(), Box<dyn Error>> {
    encoding_eq!(
        ServerMessage::Authentication(Authentication::Ok),
        b"\x52\0\0\0\x08\x00\x00\x00\x00");
    encoding_eq!(
        ServerMessage::Authentication(Authentication::Sasl {
            methods: vec![String::from("SCRAM-SHA-256")],
        }),
        b"R\0\0\0\x1d\0\0\0\n\0\0\0\x01\0\0\0\rSCRAM-SHA-256");
    encoding_eq!(
        ServerMessage::Authentication(Authentication::SaslContinue {
            data: Bytes::from_static(b"sasl_interim_data"),
        }),
        b"R\0\0\0\x1d\x00\x00\x00\x0b\0\0\0\x11sasl_interim_data");
    encoding_eq!(
        ServerMessage::Authentication(Authentication::SaslFinal {
            data: Bytes::from_static(b"sasl_final_data"),
        }),
        b"R\0\0\0\x1b\x00\x00\x00\x0c\0\0\0\x0fsasl_final_data");
    Ok(())
}

#[test]
fn log_message() -> Result<(), Box<dyn Error>> {
    encoding_eq!(
        ServerMessage::LogMessage(LogMessage {
            severity: MessageSeverity::Notice,
            code: 0xF0_00_00_00,
            text: "changing system config".into(),
            attributes: map!{},
        }),
        b"L\0\0\0%<\xf0\0\0\0\0\0\0\x16changing system config\0\0");
    Ok(())
}

#[test]
fn state_data_description() -> Result<(), Box<dyn Error>> {
    encoding_eq!(
        ServerMessage::StateDataDescription(StateDataDescription {
            typedesc: RawTypedesc {
                proto: ProtocolVersion::current(),
                id: Uuid::from_u128(0x105),
                data: Bytes::from_static(
                    b"\x02\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x01\x05"),
            },
        }),
        b"s\0\0\0)\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x01\x05\0\0\0\
        \x11\x02\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x01\x05"
    );
    Ok(())
}