diff --git a/crates/harness-tests/src/routing.rs b/crates/harness-tests/src/routing.rs index bb2c5fdca..5f1497906 100644 --- a/crates/harness-tests/src/routing.rs +++ b/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. diff --git a/crates/harness/src/monolith.rs b/crates/harness/src/monolith.rs index d15a33198..8bfb4852f 100644 --- a/crates/harness/src/monolith.rs +++ b/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, /// A mapping from request path to response body for mocking HTTP responses. response_mocks: HashMap, + rooms: HashMap, } impl Monolith { @@ -237,6 +238,35 @@ impl Monolith { pub fn collect_mock_http(&self) -> Vec { self.state.lock().unwrap().received_http.clone() } + + pub async fn load_room(&mut self, room: impl Into + 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 + 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 {