[ECS] Infra for comprehensive unit tests

crates/valence_new/src/lib.rs

@@ -11,9 +11,12 @@ pub mod instance;
 pub mod inventory;
 pub mod math;
 mod packet;
+pub(crate) mod packet_stream;
 pub mod player_list;
 pub mod player_textures;
 pub mod server;
+#[cfg(test)]
+mod unit_test;
 
 /// A [`Component`] for marking entities that should be despawned at the end of
 /// the tick.

crates/valence_new/src/packet_stream.rs

@@ -0,0 +1,181 @@
+use std::collections::VecDeque;
+
+use bytes::{BufMut, BytesMut};
+use valence_protocol::packets::S2cPlayPacket;
+use valence_protocol::{DecodePacket, EncodePacket, PacketDecoder, PacketEncoder};
+
+use crate::server::byte_channel::{ByteReceiver, ByteSender, TryRecvError};
+
+/// Represents a byte stream that packets can be read from and written to.
+pub trait PacketStream {
+    /// Parses and returns the next packet in the stream.
+    fn try_recv<'a, P>(&'a mut self) -> anyhow::Result<Option<P>>
+    where
+        P: DecodePacket<'a> + std::fmt::Debug;
+
+    /// Encodes and writes the given packet to the stream.
+    fn try_send<P>(&mut self, packet: P) -> anyhow::Result<()>
+    where
+        P: EncodePacket;
+}
+
+pub(crate) struct RealPacketStream {
+    dec: PacketDecoder,
+    recv: ByteReceiver,
+    enc: PacketEncoder,
+    send: ByteSender,
+}
+
+impl RealPacketStream {
+    pub(crate) fn new(recv: ByteReceiver, send: ByteSender) -> Self {
+        Self {
+            dec: PacketDecoder::new(),
+            recv,
+            enc: PacketEncoder::new(),
+            send,
+        }
+    }
+
+    pub fn flush(&mut self) -> anyhow::Result<()> {
+        let bytes = self.enc.take();
+        self.send.try_send(bytes)?;
+        Ok(())
+    }
+}
+
+impl PacketStream for RealPacketStream {
+    fn try_recv<'a, P>(&'a mut self) -> anyhow::Result<Option<P>>
+    where
+        P: DecodePacket<'a> + std::fmt::Debug,
+    {
+        loop {
+            match self.recv.try_recv() {
+                Ok(bytes) => self.dec.queue_bytes(bytes),
+                Err(TryRecvError::Empty) => break,
+                Err(TryRecvError::Disconnected) => anyhow::bail!("Disconnected"),
+            }
+        }
+
+        self.dec.try_next_packet()
+    }
+
+    fn try_send<P>(&mut self, packet: P) -> anyhow::Result<()>
+    where
+        P: EncodePacket,
+    {
+        self.enc.append_packet(&packet)?;
+        Ok(())
+    }
+}
+
+/// A `PacketStream` that reads and writes from an in memory buffer of packets
+/// used for testing.
+pub(crate) struct MockPacketStream {
+    recv_enc: PacketEncoder,
+    recv_dec: PacketDecoder,
+
+    send_queue: VecDeque<Vec<u8>>,
+}
+
+impl<'a> MockPacketStream {
+    #[allow(dead_code)]
+    pub(crate) fn new() -> Self {
+        Self {
+            recv_enc: PacketEncoder::new(),
+            recv_dec: PacketDecoder::new(),
+            send_queue: VecDeque::new(),
+        }
+    }
+
+    /// Injects a packet into the receive stream as if it were received from a
+    /// client.
+    ///
+    /// ```rust
+    /// use valence_new::packet_stream::MockPacketStream;
+    /// use valence_protocol::packets::c2s::play::KeepAliveC2s;
+    ///
+    /// let mut stream = MockPacketStream::new();
+    /// let packet = KeepAliveC2s { id: 0xdeadbeef };
+    /// stream.inject_recv(packet);
+    /// ```
+    #[allow(dead_code)]
+    pub(crate) fn inject_recv<P>(&mut self, packet: P)
+    where
+        P: EncodePacket,
+    {
+        self.recv_enc
+            .append_packet(&packet)
+            .expect("failed to encode injected packet");
+        let bytes = self.recv_enc.take();
+        self.recv_dec.queue_bytes(bytes);
+    }
+
+    fn queue_send<P>(&mut self, packet: P) -> anyhow::Result<()>
+    where
+        P: EncodePacket,
+    {
+        let bytes = BytesMut::new();
+        let mut w = bytes.writer();
+        P::encode_packet(&packet, &mut w)?;
+        let bytes = w.into_inner();
+        self.send_queue.push_back(bytes.to_vec());
+        Ok(())
+    }
+
+    #[allow(dead_code)]
+    pub(crate) fn flush_sent(&'a mut self) -> anyhow::Result<Vec<S2cPlayPacket<'a>>> {
+        let mut packets = Vec::new();
+        for pkt in self.send_queue.iter() {
+            let mut p: Box<&'a [u8]> = Box::new(pkt.as_slice());
+            let packet = S2cPlayPacket::<'a>::decode_packet(&mut p)?;
+            packets.push(packet);
+        }
+        Ok(packets)
+    }
+}
+
+impl PacketStream for MockPacketStream {
+    fn try_recv<'a, P>(&'a mut self) -> anyhow::Result<Option<P>>
+    where
+        P: DecodePacket<'a> + std::fmt::Debug,
+    {
+        self.recv_dec.try_next_packet()
+    }
+
+    fn try_send<P>(&mut self, packet: P) -> anyhow::Result<()>
+    where
+        P: EncodePacket,
+    {
+        self.queue_send(packet)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use valence_protocol::packets::c2s::play::KeepAliveC2s;
+    use valence_protocol::packets::s2c::play::KeepAliveS2c;
+
+    use super::*;
+
+    #[test]
+    fn test_mock_stream_read() {
+        let mut stream = MockPacketStream::new();
+        let packet = KeepAliveC2s { id: 0xdeadbeef };
+        stream.inject_recv(packet.clone());
+        let packet_out = stream.try_recv::<KeepAliveC2s>().unwrap().unwrap();
+        assert_eq!(packet.id, packet_out.id);
+    }
+
+    #[test]
+    fn test_mock_stream_assert_sent() {
+        let mut stream = MockPacketStream::new();
+        let packet = KeepAliveS2c { id: 0xdeadbeef };
+        stream.try_send(packet.clone()).unwrap();
+        let packets_out = stream.flush_sent().unwrap();
+        let S2cPlayPacket::KeepAliveS2c(packet_out) = packets_out[0] else {
+            assert!(false);
+            return;
+        };
+        assert_eq!(packet.id, packet_out.id);
+    }
+}

crates/valence_new/src/server.rs

@@ -40,7 +40,7 @@ use crate::player_list::{update_player_list, PlayerList};
 use crate::server::connect::do_accept_loop;
 use crate::Despawned;
 
-mod byte_channel;
+pub(crate) mod byte_channel;
 mod connect;
 mod packet_manager;
 
@@ -221,6 +221,12 @@ impl SharedServer {
         self.0.connection_sema.close();
         *self.0.shutdown_result.lock().unwrap() = Some(res.map_err(|e| e.into()));
     }
+
+    /// Forcefully aquires a permit to connect to the server. This is useful
+    /// for testing.
+    pub(crate) fn force_aquire_owned(&self) -> OwnedSemaphorePermit {
+        self.0.connection_sema.clone().try_acquire_owned().unwrap()
+    }
 }
 
 /// Contains information about a new client joining the server.

crates/valence_new/src/server/packet_manager.rs

@@ -195,6 +195,21 @@ pub struct PlayPacketSender {
 }
 
 impl PlayPacketSender {
+    /// Creates a new `PlayPacketSender` that will send packets such that they
+    /// can be received and asserted on by a unit test.
+    pub(crate) fn new_injectable() -> (Self, ByteReceiver) {
+        let (send, recv) = byte_channel(0);
+        (
+            PlayPacketSender {
+                enc: PacketEncoder::new(),
+                send,
+                writer_task: None,
+                handle: Handle::current(),
+            },
+            recv,
+        )
+    }
+
     pub fn append_packet<P>(&mut self, pkt: &P) -> Result<()>
     where
         P: EncodePacket + ?Sized,
@@ -262,6 +277,20 @@ pub struct PlayPacketReceiver {
 }
 
 impl PlayPacketReceiver {
+    /// Creates a new `PlayPacketReceiver` with a new `ByteSender` to manually
+    /// inject packets. This is useful for testing.
+    pub(crate) fn new_injectable() -> (Self, ByteSender) {
+        let (send, recv) = byte_channel(8192);
+        (
+            Self {
+                dec: Default::default(),
+                recv,
+                reader_task: tokio::spawn(async {}),
+            },
+            send,
+        )
+    }
+
     pub fn try_next_packet<'a, P>(&'a mut self) -> Result<Option<P>>
     where
         P: DecodePacket<'a> + fmt::Debug,

crates/valence_new/src/unit_test/example.rs

@@ -0,0 +1,61 @@
+use bevy_app::App;
+use bevy_ecs::prelude::*;
+
+use super::util::create_mock_client;
+use crate::config::ServerPlugin;
+use crate::server::Server;
+use crate::unit_test::util::gen_client_info;
+
+/// Examples of valence unit tests that need to test the behavior of the server,
+/// and not just the logic of a single function. This module is meant to be a
+/// pallette of examples for how to write such tests, with various levels of
+/// complexity.
+///
+/// Some of the tests in this file may be inferior duplicates of real tests.
+#[cfg(test)]
+mod tests {
+    use bytes::BytesMut;
+
+    use super::*;
+    use crate::client::Client;
+
+    /// The server's tick should increment every update.
+    #[test]
+    fn test_server_tick_increment() {
+        let mut app = App::new();
+        app.add_plugin(ServerPlugin::new(()));
+        let server = app.world.resource::<Server>();
+        let tick = server.current_tick();
+        drop(server);
+        app.update();
+        let server = app.world.resource::<Server>();
+        assert_eq!(server.current_tick(), tick + 1);
+    }
+
+    /// A unit test where we want to test what happens when a client sends a
+    /// packet to the server.
+    #[tokio::test]
+    async fn test_client_position() {
+        let mut app = App::new();
+        app.add_plugin(ServerPlugin::new(()));
+        let server = app.world.resource::<Server>();
+        let permit = server.force_aquire_owned();
+        let info = gen_client_info("test");
+        let (client, mut client_helper) = create_mock_client(permit, info);
+        let client_ent = app.world.spawn(client).id();
+
+        // Send a packet as the client to the server.
+        let packet = valence_protocol::packets::c2s::play::SetPlayerPosition {
+            position: [12.0, 64.0, 0.0],
+            on_ground: true,
+        };
+        client_helper.send_packet(packet);
+
+        // Process the packet.
+        app.update();
+
+        // Make assertions
+        let client: &Client = app.world.get(client_ent).unwrap();
+        assert_eq!(client.position(), [12.0, 64.0, 0.0].into());
+    }
+}

crates/valence_new/src/unit_test/mod.rs

@@ -0,0 +1,2 @@
+mod example;
+pub(crate) mod util;

crates/valence_new/src/unit_test/util.rs

@@ -0,0 +1,59 @@
+use bytes::BytesMut;
+use tokio::sync::OwnedSemaphorePermit;
+use valence_protocol::{EncodePacket, Username};
+
+use crate::client::Client;
+use crate::server::byte_channel::{ByteReceiver, ByteSender};
+use crate::server::{NewClientInfo, PlayPacketReceiver, PlayPacketSender};
+
+/// Creates a mock client that can be used for unit testing.
+///
+/// Returns the client, and a helper to inject packets as if the client sent
+/// them and receive packets as if the client received them.
+pub fn create_mock_client(
+    permit: OwnedSemaphorePermit,
+    client_info: NewClientInfo,
+) -> (Client, MockClientHelper) {
+    let (pkt_send, recv) = PlayPacketSender::new_injectable();
+    let (pkt_recv, send) = PlayPacketReceiver::new_injectable();
+
+    let client = Client::new(pkt_send, pkt_recv, permit, client_info);
+    (client, MockClientHelper::new(send, recv))
+}
+
+/// Creates a `NewClientInfo` with the given username and a random UUID.
+/// Panics if the username is invalid.
+pub fn gen_client_info(username: &str) -> NewClientInfo {
+    NewClientInfo {
+        username: Username::new(username.to_owned()).unwrap(),
+        uuid: uuid::Uuid::new_v4(),
+        ip: std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1)),
+        properties: vec![],
+    }
+}
+
+pub struct MockClientHelper {
+    /// Used to pretend the client sent bytes.
+    send: ByteSender,
+    /// Used to pretend the client received bytes.
+    recv: ByteReceiver,
+}
+
+/// Has a `ByteSender` to inject packets to be read and acted
+/// upon by the server, and a `ByteReceiver` to receive packets and make
+/// assertions about what the server sent.
+impl MockClientHelper {
+    fn new(send: ByteSender, recv: ByteReceiver) -> Self {
+        Self { send, recv }
+    }
+
+    /// Inject a packet to be parsed by the server. Panics if the packet cannot
+    /// be sent.
+    pub fn send_packet(&mut self, packet: impl EncodePacket) {
+        let mut buffer = Vec::<u8>::new();
+        valence_protocol::encode_packet(&mut buffer, &packet).expect("Failed to encode packet");
+        self.send
+            .try_send(BytesMut::from(buffer.as_slice()))
+            .expect("Failed to send packet");
+    }
+}