add un/load_room to emulated monoliths (#1105)

* add `un`/`load_room` to emulated monoliths

* fix lints

crates/harness-tests/src/routing.rs

@@ -1,7 +1,7 @@
 use std::time::Duration;
 
 use harness::{Client, MockRespParts, Monolith, TestRunner};
-use ott_balancer_protocol::monolith::{MsgM2B, RoomMetadata};
+
 use test_context::test_context;
 
 #[test_context(TestRunner)]
@@ -20,11 +20,7 @@ async fn route_http_to_correct_monolith(ctx: &mut TestRunner) {
     );
 
     m.show().await;
-    m.send(MsgM2B::Loaded {
-        name: "foo".to_owned().into(),
-        metadata: RoomMetadata::default(),
-    })
-    .await;
+    m.load_room("foo".to_string()).await;
 
     // Without this sleep, this test can trigger a race condition where the client connects to the balancer before the monolith has the room loaded.
     // This will cause the other monolith to get the room loaded, and the client will connect to that monolith instead.
@@ -104,11 +100,7 @@ async fn route_ws_to_correct_monolith(ctx: &mut TestRunner) {
 
     let mut m = Monolith::new(ctx).await.unwrap();
     m.show().await;
-    m.send(MsgM2B::Loaded {
-        name: "foo".to_owned().into(),
-        metadata: RoomMetadata::default(),
-    })
-    .await;
+    m.load_room("foo".to_string()).await;
 
     // Without this sleep, this test can trigger a race condition where the client connects to the balancer before the monolith has the room loaded.
     // This will cause the other monolith to get the room loaded, and the client will connect to that monolith instead.

crates/harness/src/monolith.rs

@@ -11,7 +11,7 @@ use http_body_util::{BodyExt, Full};
 use hyper::{body::Incoming as IncomingBody, Request};
 use hyper::{service::Service, Response};
 
-use ott_balancer_protocol::monolith::*;
+use ott_balancer_protocol::{monolith::*, RoomName};
 use tokio::{net::TcpListener, sync::Notify};
 use tracing::warn;
 use tungstenite::Message;
@@ -47,6 +47,7 @@ pub(crate) struct MonolithState {
     received_http: Vec<MockRequest>,
     /// A mapping from request path to response body for mocking HTTP responses.
     response_mocks: HashMap<String, (MockRespParts, Bytes)>,
+    rooms: HashMap<RoomName, RoomMetadata>,
 }
 
 impl Monolith {
@@ -237,6 +238,35 @@ impl Monolith {
     pub fn collect_mock_http(&self) -> Vec<MockRequest> {
         self.state.lock().unwrap().received_http.clone()
     }
+
+    pub async fn load_room(&mut self, room: impl Into<RoomName> + Clone) {
+        let room = room.into();
+        let meta = RoomMetadata::default();
+        self.state
+            .lock()
+            .unwrap()
+            .rooms
+            .insert(room.clone(), meta.clone());
+        if self.connected() {
+            self.send(MsgM2B::Loaded {
+                name: room,
+                metadata: meta,
+            })
+            .await;
+        }
+    }
+
+    pub async fn unload_room(&mut self, room: impl Into<RoomName> + Clone) {
+        let room = room.into();
+        self.state
+            .lock()
+            .unwrap()
+            .rooms
+            .remove(&room.clone().clone());
+        if self.connected() {
+            self.send(MsgM2B::Unloaded { room }).await;
+        }
+    }
 }
 
 impl Drop for Monolith {