wsframe tests

src/wsframe.rs

@@ -124,10 +124,10 @@ impl Frame {
             let mut length = u64::from(second & 0x7F);
 
             if length == 126 {
-                let mut length_bytes = [0u8; 2];
                 if size < 2 {
                     return Ok(None)
                 }
+                let mut length_bytes = [0u8; 2];
                 length_bytes.copy_from_slice(&buf[idx..idx+2]);
                 size -= 2;
                 idx += 2;
@@ -138,13 +138,13 @@ impl Frame {
                     wide});
                 header_length += 2;
             } else if length == 127 {
-                let mut length_bytes = [0u8; 8];
                 if size < 8 {
                     return Ok(None)
                 }
+                let mut length_bytes = [0u8; 8];
                 length_bytes.copy_from_slice(&buf[idx..idx+8]);
                 size -= 8;
-                idx += 2;
+                idx += 8;
 
                 unsafe { length = mem::transmute(length_bytes); }
                 length = u64::from_be(length);
@@ -329,3 +329,100 @@ impl fmt::Display for Frame {
             self.payload.iter().map(|byte| format!("{:x}", byte)).collect::<String>())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse() {
+        let mut buf = BytesMut::from(&[0b00000001u8, 0b00000001u8][..]);
+        assert!(Frame::parse(&mut buf).unwrap().is_none());
+        buf.extend(b"1");
+        let frame = Frame::parse(&mut buf).unwrap().unwrap();
+        println!("FRAME: {:?}", frame);
+        assert!(!frame.finished);
+        assert_eq!(frame.opcode, OpCode::Text);
+        assert_eq!(frame.payload, &b"1"[..]);
+    }
+
+    #[test]
+    fn test_parse_length0() {
+        let mut buf = BytesMut::from(&[0b00000001u8, 0b00000000u8][..]);
+        let frame = Frame::parse(&mut buf).unwrap().unwrap();
+        assert!(!frame.finished);
+        assert_eq!(frame.opcode, OpCode::Text);
+        assert!(frame.payload.is_empty());
+    }
+
+    #[test]
+    fn test_parse_length2() {
+        let mut buf = BytesMut::from(&[0b00000001u8, 126u8][..]);
+        assert!(Frame::parse(&mut buf).unwrap().is_none());
+        buf.extend(&[0u8, 4u8][..]);
+        buf.extend(b"1234");
+
+        let frame = Frame::parse(&mut buf).unwrap().unwrap();
+        assert!(!frame.finished);
+        assert_eq!(frame.opcode, OpCode::Text);
+        assert_eq!(frame.payload, &b"1234"[..]);
+    }
+
+    #[test]
+    fn test_parse_length4() {
+        let mut buf = BytesMut::from(&[0b00000001u8, 127u8][..]);
+        assert!(Frame::parse(&mut buf).unwrap().is_none());
+        buf.extend(&[0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 4u8][..]);
+        buf.extend(b"1234");
+
+        let frame = Frame::parse(&mut buf).unwrap().unwrap();
+        assert!(!frame.finished);
+        assert_eq!(frame.opcode, OpCode::Text);
+        assert_eq!(frame.payload, &b"1234"[..]);
+    }
+
+    #[test]
+    fn test_parse_frame_mask() {
+        let mut buf = BytesMut::from(&[0b00000001u8, 0b10000001u8][..]);
+        buf.extend(b"0001");
+        buf.extend(b"1");
+
+        let frame = Frame::parse(&mut buf).unwrap().unwrap();
+        assert!(!frame.finished);
+        assert_eq!(frame.opcode, OpCode::Text);
+        assert_eq!(frame.payload, vec![1u8]);
+    }
+
+    #[test]
+    fn test_ping_frame() {
+        let mut frame = Frame::message(Vec::from("data"), OpCode::Ping, true);
+        let mut buf = Vec::new();
+        frame.format(&mut buf).unwrap();
+
+        let mut v = vec![137u8, 4u8];
+        v.extend(b"data");
+        assert_eq!(buf, v);
+    }
+
+    #[test]
+    fn test_pong_frame() {
+        let mut frame = Frame::message(Vec::from("data"), OpCode::Pong, true);
+        let mut buf = Vec::new();
+        frame.format(&mut buf).unwrap();
+
+        let mut v = vec![138u8, 4u8];
+        v.extend(b"data");
+        assert_eq!(buf, v);
+    }
+
+    #[test]
+    fn test_close_frame() {
+        let mut frame = Frame::close(CloseCode::Normal, "data");
+        let mut buf = Vec::new();
+        frame.format(&mut buf).unwrap();
+
+        let mut v = vec![136u8, 6u8, 3u8, 232u8];
+        v.extend(b"data");
+        assert_eq!(buf, v);
+    }
+}