chore(engine): roxiproxi tests

Author: NathanFlurry

Cargo.lock

@@ -8,6 +8,7 @@ pub struct TestOpts {
 	pub timeout_secs: u64,
 	pub pegboard_outbound: bool,
 	pub auth_admin_token: Option<String>,
+	pub network_faults: bool,
 }
 
 impl TestOpts {
@@ -17,6 +18,7 @@ impl TestOpts {
 			timeout_secs: 10,
 			pegboard_outbound: false,
 			auth_admin_token: None,
+			network_faults: false,
 		}
 	}
 
@@ -34,6 +36,11 @@ impl TestOpts {
 		self.auth_admin_token = Some(token.into());
 		self
 	}
+
+	pub fn with_network_faults(mut self) -> Self {
+		self.network_faults = true;
+		self
+	}
 }
 
 impl Default for TestOpts {
@@ -43,13 +50,15 @@ impl Default for TestOpts {
 			timeout_secs: 10,
 			pegboard_outbound: false,
 			auth_admin_token: None,
+			network_faults: false,
 		}
 	}
 }
 
 pub struct TestCtx {
 	dcs: Vec<TestDatacenter>,
 	pub opts: TestOpts,
+	network_faults: Option<rivet_test_deps::ToxiproxyTestServer>,
 }
 
 pub struct TestDatacenter {
@@ -99,7 +108,17 @@ impl TestCtx {
 			futures_util::future::try_join_all(setup_futures).await?;
 		dcs.sort_by_key(|dc| dc.config.dc_label());
 
-		Ok(Self { dcs, opts })
+		let network_faults = if opts.network_faults {
+			Some(rivet_test_deps::ToxiproxyTestServer::start().await?)
+		} else {
+			None
+		};
+
+		Ok(Self {
+			dcs,
+			opts,
+			network_faults,
+		})
 	}
 
 	async fn setup_instance(
@@ -204,6 +223,12 @@ impl TestCtx {
 			.unwrap_or_else(|| panic!("No datacenter found with label {}", label))
 	}
 
+	pub fn network_faults(&self) -> &rivet_test_deps::ToxiproxyTestServer {
+		self.network_faults
+			.as_ref()
+			.expect("Network faults were not enabled. Use TestOpts::with_network_faults().")
+	}
+
 	pub async fn shutdown(self) {
 		tracing::info!("shutting down multi-DC test context");
 		for dc in self.dcs {

engine/packages/engine/tests/common/ctx.rs

@@ -20,6 +20,46 @@ type ActorFactory = Arc<dyn Fn(ActorConfig) -> Box<dyn TestActor> + Send + Sync>
 
 pub type TestEnvoy = Envoy;
 
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum EnvoyConnectionEvent {
+	Connected,
+	Disconnected,
+}
+
+pub struct EnvoyConnectionEventWaiter {
+	rx: broadcast::Receiver<EnvoyConnectionEvent>,
+	expected: EnvoyConnectionEvent,
+	timeout: std::time::Duration,
+}
+
+impl EnvoyConnectionEventWaiter {
+	pub fn assert_no_event(&mut self) {
+		match self.rx.try_recv() {
+			Err(tokio::sync::broadcast::error::TryRecvError::Empty) => {}
+			Ok(event) => panic!("unexpected Envoy connection event before fault: {event:?}"),
+			Err(tokio::sync::broadcast::error::TryRecvError::Lagged(count)) => {
+				panic!("missed {count} Envoy connection events before fault")
+			}
+			Err(err) => panic!("Envoy connection event channel closed: {err}"),
+		}
+	}
+
+	pub async fn wait(mut self) {
+		tokio::time::timeout(self.timeout, async {
+			loop {
+				match self.rx.recv().await {
+					Ok(event) if event == self.expected => break,
+					Ok(_) => {}
+					Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {}
+					Err(err) => panic!("Envoy connection event channel closed: {err}"),
+				}
+			}
+		})
+		.await
+		.expect("timed out waiting for Envoy connection event");
+	}
+}
+
 #[derive(Clone)]
 pub struct EnvoyConfig {
 	endpoint: String,
@@ -119,6 +159,7 @@ impl EnvoyBuilder {
 				actor_factories: self.actor_factories,
 				actors: tokio::sync::Mutex::new(HashMap::new()),
 				lifecycle_tx,
+				connection_tx: broadcast::channel(100).0,
 			}),
 			handle: tokio::sync::Mutex::new(None),
 			envoy_key: uuid::Uuid::new_v4().to_string(),
@@ -130,6 +171,7 @@ struct EnvoyInner {
 	actor_factories: HashMap<String, ActorFactory>,
 	actors: tokio::sync::Mutex<HashMap<String, Box<dyn TestActor>>>,
 	lifecycle_tx: broadcast::Sender<ActorLifecycleEvent>,
+	connection_tx: broadcast::Sender<EnvoyConnectionEvent>,
 }
 
 pub struct Envoy {
@@ -197,6 +239,21 @@ impl Envoy {
 		self.inner.lifecycle_tx.subscribe()
 	}
 
+	pub fn subscribe_connection_events(&self) -> broadcast::Receiver<EnvoyConnectionEvent> {
+		self.inner.connection_tx.subscribe()
+	}
+
+	pub fn wait_for_next_connection_event(
+		&self,
+		expected: EnvoyConnectionEvent,
+	) -> EnvoyConnectionEventWaiter {
+		EnvoyConnectionEventWaiter {
+			rx: self.subscribe_connection_events(),
+			expected,
+			timeout: std::time::Duration::from_secs(20),
+		}
+	}
+
 	pub async fn shutdown(&self) {
 		if let Some(handle) = self.handle.lock().await.take() {
 			handle.shutdown_and_wait(false).await;
@@ -241,6 +298,17 @@ impl TestEnvoyCallbacks {
 }
 
 impl rivet_test_envoy::EnvoyCallbacks for TestEnvoyCallbacks {
+	fn on_connect(&self, _handle: EnvoyHandle) {
+		let _ = self.inner.connection_tx.send(EnvoyConnectionEvent::Connected);
+	}
+
+	fn on_disconnect(&self, _handle: EnvoyHandle) {
+		let _ = self
+			.inner
+			.connection_tx
+			.send(EnvoyConnectionEvent::Disconnected);
+	}
+
 	fn on_actor_start(
 		&self,
 		handle: EnvoyHandle,
@@ -527,6 +595,7 @@ pub struct TestEnvoyBuilder {
 	namespace: String,
 	pool_name: String,
 	version: u32,
+	endpoint: Option<String>,
 	actor_factories: HashMap<String, ActorFactory>,
 }
 
@@ -536,6 +605,7 @@ impl TestEnvoyBuilder {
 			namespace: namespace.to_string(),
 			pool_name: "test-envoy".to_string(),
 			version: 1,
+			endpoint: None,
 			actor_factories: HashMap::new(),
 		}
 	}
@@ -550,6 +620,11 @@ impl TestEnvoyBuilder {
 		self
 	}
 
+	pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
+		self.endpoint = Some(endpoint.into());
+		self
+	}
+
 	pub fn with_actor_behavior<F>(mut self, actor_name: &str, factory: F) -> Self
 	where
 		F: Fn(ActorConfig) -> Box<dyn TestActor> + Send + Sync + 'static,
@@ -561,7 +636,10 @@ impl TestEnvoyBuilder {
 
 	pub async fn build(self, dc: &super::TestDatacenter) -> Result<Envoy> {
 		let config = EnvoyConfig::builder()
-			.endpoint(format!("http://127.0.0.1:{}", dc.guard_port()))
+			.endpoint(
+				self.endpoint
+					.unwrap_or_else(|| format!("http://127.0.0.1:{}", dc.guard_port())),
+			)
 			.token("dev")
 			.namespace(&self.namespace)
 			.pool_name(&self.pool_name)

engine/packages/engine/tests/common/test_envoy.rs

@@ -1,9 +1,31 @@
-use std::{collections::HashMap, str::FromStr};
+use std::{collections::HashMap, future::Future, str::FromStr};
 
 use serde_json::json;
 
 use super::TestDatacenter;
 
+pub async fn wait_with_poll<T, F, Fut>(
+	timeout: std::time::Duration,
+	poll_interval: std::time::Duration,
+	mut check: F,
+) -> Option<T>
+where
+	F: FnMut() -> Fut,
+	Fut: Future<Output = Option<T>>,
+{
+	tokio::time::timeout(timeout, async {
+		loop {
+			if let Some(value) = check().await {
+				break value;
+			}
+
+			tokio::time::sleep(poll_interval).await;
+		}
+	})
+	.await
+	.ok()
+}
+
 // Namespace helpers
 pub async fn setup_test_namespace(leader_dc: &TestDatacenter) -> (String, rivet_util::Id) {
 	let random_suffix = rand::random::<u16>();

engine/packages/engine/tests/common/test_helpers.rs

@@ -11,3 +11,5 @@ pub mod api_actors_get_or_create;
 pub mod api_actors_list;
 pub mod api_actors_list_names;
 pub mod auth;
+pub mod auth;
+pub mod network_faults;