From d1f4161ee9d8c55b7e9f7c42b28a9b3277470032 Mon Sep 17 00:00:00 2001
From: cloudhead <cloudhead@radicle.xyz>
Date: Fri, 9 Feb 2024 11:31:02 +0100
Subject: [PATCH] sim: Track messages received in simulator

This is going to be useful for testing message amplification in the
network protocol.

We also reset the state of a few things everytime the simulator is run.
---
 radicle-node/src/test/simulator.rs | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/radicle-node/src/test/simulator.rs b/radicle-node/src/test/simulator.rs
index 26ff9bec..859d6894 100644
--- a/radicle-node/src/test/simulator.rs
+++ b/radicle-node/src/test/simulator.rs
@@ -63,7 +63,7 @@ pub enum Input {
     },
     /// Disconnected from peer.
     Disconnected(NodeId, Rc<DisconnectReason>),
-    /// Received a message from a remote peer.
+    /// Received messages from a remote peer.
     Received(NodeId, Vec<Message>),
     /// Fetch completed for a node.
     Fetched(RepoId, NodeId, Rc<Result<fetch::FetchResult, FetchError>>),
@@ -175,6 +175,8 @@ pub struct Simulation<S, G> {
     inbox: Inbox,
     /// Events emitted during simulation.
     events: BTreeMap<NodeId, VecDeque<Event>>,
+    /// Messages received during simulation.
+    messages: Vec<Message>,
     /// Priority events that should happen immediately.
     priority: VecDeque<Scheduled>,
     /// Simulated latencies between nodes.
@@ -207,6 +209,7 @@ impl<S: WriteStorage + 'static, G: Signer> Simulation<S, G> {
                 messages: BTreeMap::new(),
             },
             events: BTreeMap::new(),
+            messages: Vec::new(),
             priority: VecDeque::new(),
             partitions: BTreeSet::new(),
             latencies: BTreeMap::new(),
@@ -246,6 +249,11 @@ impl<S: WriteStorage + 'static, G: Signer> Simulation<S, G> {
         self.events.entry(*node).or_default().drain(..)
     }
 
+    /// Get all messages received by nodes during the simulation.
+    pub fn messages(&mut self) -> &[Message] {
+        &self.messages
+    }
+
     /// Get the latency between two nodes. The minimum latency between nodes is 1 millisecond.
     pub fn latency(&self, from: NodeId, to: NodeId) -> LocalDuration {
         self.latencies
@@ -294,6 +302,10 @@ impl<S: WriteStorage + 'static, G: Signer> Simulation<S, G> {
     {
         let mut nodes: BTreeMap<_, _> = peers.into_iter().map(|p| (p.id(), p)).collect();
 
+        self.messages.clear();
+        self.events.clear();
+        self.start_time = self.time;
+
         while self.step_(&mut nodes) {
             if !pred(self) {
                 break;
@@ -406,9 +418,10 @@ impl<S: WriteStorage + 'static, G: Signer> Simulation<S, G> {
                     }
                     Input::Wake => p.wake(),
                     Input::Received(id, msgs) => {
-                        for msg in msgs {
+                        for msg in msgs.clone() {
                             p.received_message(id, msg);
                         }
+                        self.messages.extend(msgs);
                     }
                     Input::Fetched(rid, nid, result) => {
                         let result = Rc::try_unwrap(result).unwrap();