@rivetkit/engine 1.0.0 → 2.2.1-pr.4600.b74ff3b

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

@@ -1,42 +1,55 @@
-use std::time::Duration;
+use std::{collections::HashMap, str::FromStr};
 
 use serde_json::json;
 
+use super::TestDatacenter;
+
 // Namespace helpers
-pub async fn setup_test_namespace(guard_port: u16) -> (String, rivet_util::Id) {
+pub async fn setup_test_namespace(leader_dc: &TestDatacenter) -> (String, rivet_util::Id) {
 	let random_suffix = rand::random::<u16>();
 	let namespace_name = format!("test-{random_suffix}");
-	let namespace_id = super::create_namespace(&namespace_name, guard_port).await;
-	(namespace_name, namespace_id)
+	let res = super::api::public::namespaces_create(
+		leader_dc.guard_port(),
+		rivet_api_peer::namespaces::CreateRequest {
+			name: namespace_name,
+			display_name: "Test Namespace".to_string(),
+		},
+	)
+	.await
+	.expect("failed to setup test namespace");
+	(res.namespace.name, res.namespace.namespace_id)
 }
 
 // Setup namespace with runner
 pub async fn setup_test_namespace_with_runner(
 	dc: &super::TestDatacenter,
-) -> (String, rivet_util::Id, super::runner::TestRunner) {
-	let (namespace_name, namespace_id) = setup_test_namespace(dc.guard_port()).await;
+) -> (String, rivet_util::Id, super::test_runner::TestRunner) {
+	let (namespace_name, namespace_id) = setup_test_namespace(dc).await;
 
-	let runner = super::runner::TestRunner::new(
-		dc.guard_port(),
-		&namespace_name,
-		&format!("key-{:012x}", rand::random::<u64>()),
-		1,
-		20,
-	)
+	let runner = setup_runner(dc, &namespace_name, |builder| {
+		builder.with_actor_behavior("test-actor", |_config| {
+			Box::new(super::test_runner::EchoActor::new())
+		})
+	})
 	.await;
 
 	(namespace_name, namespace_id, runner)
 }
 
-pub async fn setup_runner(
+// Setup namespace with envoy
+pub async fn setup_test_namespace_with_envoy(
 	dc: &super::TestDatacenter,
-	namespace_name: &str,
-	key: &str,
-	version: u32,
-	total_slots: u32,
-) -> super::runner::TestRunner {
-	super::runner::TestRunner::new(dc.guard_port(), &namespace_name, key, version, total_slots)
-		.await
+) -> (String, rivet_util::Id, super::test_envoy::TestEnvoy) {
+	let (namespace_name, namespace_id) = setup_test_namespace(dc).await;
+
+	let envoy = setup_envoy(dc, &namespace_name, |builder| {
+		builder.with_actor_behavior("test-actor", |_config| {
+			Box::new(super::test_envoy::EchoActor::new())
+		})
+	})
+	.await;
+
+	(namespace_name, namespace_id, envoy)
 }
 
 pub async fn cleanup_test_namespace(namespace_id: rivet_util::Id, _guard_port: u16) {
@@ -71,77 +84,45 @@ pub fn generate_special_chars_string() -> String {
 	"test-!@#$%^&*()_+-=[]{}|;':\",./<>?".to_string()
 }
 
-// Wait helpers
-pub async fn wait_for_actor_propagation(actor_id: &str, timeout_secs: u64) {
-	tracing::info!(?actor_id, ?timeout_secs, "waiting for actor propagation");
-	tokio::time::sleep(Duration::from_secs(timeout_secs)).await;
-}
-
-pub async fn wait_for_eventual_consistency() {
-	tracing::info!("waiting for eventual consistency");
-	tokio::time::sleep(Duration::from_millis(500)).await;
-}
-
 // Actor verification helpers
 pub async fn assert_actor_exists(
+	port: u16,
 	actor_id: &str,
 	namespace: &str,
-	guard_port: u16,
-) -> serde_json::Value {
-	let response = super::get_actor(actor_id, Some(namespace), guard_port).await;
-	assert!(
-		response.status().is_success(),
-		"Actor {} should exist in namespace {}",
-		actor_id,
-		namespace
-	);
-	response
-		.json()
-		.await
-		.expect("Failed to parse actor response")
+) -> rivet_types::actors::Actor {
+	let res = super::try_get_actor(port, actor_id, namespace).await;
+	let Some(actor) = res.expect("Failed to try_get_actor") else {
+		panic!("Actor {} should exist in namespace {}", actor_id, namespace,);
+	};
+	actor
 }
 
-pub async fn assert_actor_not_exists(actor_id: &str, guard_port: u16) {
-	let response = super::get_actor(actor_id, None, guard_port).await;
-	assert_eq!(
-		response.status(),
-		400,
-		"Actor {} should not exist (expecting 400 for Actor::NotFound)",
-		actor_id
-	);
+pub async fn assert_actor_not_exists(port: u16, actor_id: &str, namespace: &str) {
+	let res = super::try_get_actor(port, actor_id, namespace).await;
+	if let Some(actor) = res.expect("Failed to try_get_actor") {
+		panic!(
+			"Actor {} should not exist in namespace {}",
+			actor_id, namespace,
+		);
+	};
 }
 
-pub async fn assert_actor_returns_bad_request(actor_id: &str, guard_port: u16) {
-	let response = super::get_actor(actor_id, None, guard_port).await;
-	assert_eq!(
-		response.status(),
-		400,
-		"Actor {} should return 400 BAD_REQUEST",
+pub async fn assert_actor_is_destroyed(port: u16, actor_id: &str, namespace: &str) {
+	let actor = assert_actor_exists(port, actor_id, namespace).await;
+	assert!(
+		actor.destroy_ts.is_some(),
+		"Actor {} should have destroy_ts set",
 		actor_id
 	);
 }
 
-pub async fn assert_actor_is_destroyed(actor_id: &str, namespace: Option<&str>, guard_port: u16) {
-	let response = super::get_actor(actor_id, namespace, guard_port).await;
+pub async fn assert_actor_is_alive(port: u16, actor_id: &str, namespace: &str) {
+	let actor = assert_actor_exists(port, actor_id, namespace).await;
 	assert!(
-		response.status().is_success(),
-		"Actor {} should still be retrievable after destroy",
+		actor.destroy_ts.is_none(),
+		"Actor {} should not have destroy_ts set",
 		actor_id
 	);
-
-	let body: serde_json::Value = response
-		.json()
-		.await
-		.expect("Failed to parse actor response");
-
-	tracing::info!(?body, ?actor_id, "assert_actor_is_destroyed response body");
-
-	assert!(
-		body["actor"]["destroy_ts"].as_i64().is_some(),
-		"Actor {} should have destroy_ts set. Response: {:?}",
-		actor_id,
-		body
-	);
 }
 
 pub async fn assert_actor_in_dc(actor_id_str: &str, expected_dc_label: u16) {
@@ -149,6 +130,7 @@ pub async fn assert_actor_in_dc(actor_id_str: &str, expected_dc_label: u16) {
 	let actor_id: rivet_util::Id = actor_id_str.parse().expect("Failed to parse actor ID");
 	let actual_dc_label = actor_id.label();
 
+	// TODO: make this fetch as well
 	assert_eq!(
 		actual_dc_label, expected_dc_label,
 		"Actor should be in datacenter {} but is in {}",
@@ -156,71 +138,177 @@ pub async fn assert_actor_in_dc(actor_id_str: &str, expected_dc_label: u16) {
 	);
 }
 
-pub async fn assert_actor_in_runner(
-	dc: &super::TestDatacenter,
-	actor_id_str: &str,
-	expected_runner_id: &str,
+pub fn assert_actors_equal(
+	actor1: &rivet_types::actors::Actor,
+	actor2: &rivet_types::actors::Actor,
 ) {
-	let actor_id: rivet_util::Id = actor_id_str.parse().expect("Failed to parse actor ID");
-
-	let actors_res = dc
-		.workflow_ctx
-		.op(pegboard::ops::actor::get_runner::Input {
-			actor_ids: vec![actor_id],
-		})
-		.await
-		.unwrap();
-	let runner_id = actors_res.actors.first().map(|x| x.runner_id.to_string());
-
+	assert_eq!(actor1.actor_id, actor2.actor_id, "Actor IDs should match");
 	assert_eq!(
-		runner_id,
-		Some(expected_runner_id.to_string()),
-		"Actor {actor_id} should be in runner {expected_runner_id} (actually in runner {runner_id:?})",
+		actor1.namespace_id, actor2.namespace_id,
+		"Namespace IDs should match"
 	);
+	assert_eq!(actor1.name, actor2.name, "Actor names should match");
 }
 
-pub fn assert_actors_equal(actor1: &serde_json::Value, actor2: &serde_json::Value) {
-	assert_eq!(
-		actor1["actor"]["actor_id"], actor2["actor"]["actor_id"],
-		"Actor IDs should match"
-	);
-	assert_eq!(
-		actor1["actor"]["namespace_id"], actor2["actor"]["namespace_id"],
-		"Namespace IDs should match"
-	);
-	assert_eq!(
-		actor1["actor"]["name"], actor2["actor"]["name"],
-		"Actor names should match"
-	);
+// Datacenter helpers
+pub fn get_test_datacenter_name(label: u16) -> String {
+	format!("dc-{}", label)
 }
 
-// Response assertion helpers
-pub async fn assert_created_response(response: &serde_json::Value, expected_created: bool) {
-	let created = response["created"]
-		.as_bool()
-		.expect("Missing created field in response");
-	assert_eq!(
-		created, expected_created,
-		"Expected created to be {}",
-		expected_created
-	);
+// Timing helpers for eventual consistency
+pub async fn wait_for_eventual_consistency() {
+	tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
 }
 
-pub fn assert_pagination_response(response: &serde_json::Value) {
-	assert!(
-		response.get("actors").is_some() || response.get("names").is_some(),
-		"Response should have actors or names array"
+pub async fn wait_for_actor_propagation(actor_id: &str, _generation: u32) {
+	// Wait for actor state to propagate through all systems
+	tokio::time::sleep(tokio::time::Duration::from_millis(1000)).await;
+}
+
+// Test runner helper functions
+
+/// Build a test runner with specified configuration
+///
+/// Defaults to 20 total slots, but can be overridden in the builder closure.
+///
+/// Example:
+/// ```
+/// // Default 20 slots
+/// let runner = setup_runner(ctx.leader_dc(), &namespace, |builder| {
+///     builder
+///         .with_actor_behavior("test-actor", |_| Box::new(EchoActor::new()))
+///         .with_actor_behavior("crash-actor", |_| Box::new(CrashOnStartActor::new(1)))
+/// }).await;
+///
+/// // Override slots
+/// let runner = setup_runner(ctx.leader_dc(), &namespace, |builder| {
+///     builder
+///         .with_total_slots(2)
+///         .with_actor_behavior("test-actor", |_| Box::new(EchoActor::new()))
+/// }).await;
+/// ```
+pub async fn setup_runner<F>(
+	dc: &super::TestDatacenter,
+	namespace: &str,
+	configure: F,
+) -> super::test_runner::TestRunner
+where
+	F: FnOnce(super::test_runner::TestRunnerBuilder) -> super::test_runner::TestRunnerBuilder,
+{
+	let builder = super::test_runner::TestRunnerBuilder::new(namespace)
+		.with_runner_key(&format!("key-{:012x}", rand::random::<u64>()))
+		.with_version(1)
+		.with_total_slots(20);
+
+	let builder = configure(builder);
+
+	let runner = builder
+		.build(dc)
+		.await
+		.expect("failed to build test runner");
+
+	runner.start().await.expect("failed to start runner");
+	runner.wait_ready().await;
+
+	runner
+}
+
+/// Build a test envoy with specified configuration
+///
+/// Defaults to 20 total slots, but can be overridden in the builder closure.
+///
+/// Example:
+/// ```
+/// // Default 20 slots
+/// let envoy = setup_envoy(ctx.leader_dc(), &namespace, |builder| {
+///     builder
+///         .with_actor_behavior("test-actor", |_| Box::new(EchoActor::new()))
+///         .with_actor_behavior("crash-actor", |_| Box::new(CrashOnStartActor::new(1)))
+/// }).await;
+///
+/// // Override slots
+/// let envoy = setup_envoy(ctx.leader_dc(), &namespace, |builder| {
+///     builder
+///         .with_total_slots(2)
+///         .with_actor_behavior("test-actor", |_| Box::new(EchoActor::new()))
+/// }).await;
+/// ```
+pub async fn setup_envoy<F>(
+	dc: &super::TestDatacenter,
+	namespace: &str,
+	configure: F,
+) -> super::test_envoy::TestEnvoy
+where
+	F: FnOnce(super::test_envoy::TestEnvoyBuilder) -> super::test_envoy::TestEnvoyBuilder,
+{
+	let builder = super::test_envoy::TestEnvoyBuilder::new(namespace).with_version(1);
+
+	let builder = configure(builder);
+
+	let envoy = builder.build(dc).await.expect("failed to build test envoy");
+
+	// Upsert serverful runner config
+	let mut datacenters = HashMap::new();
+	datacenters.insert(
+		dc.config.dc_name().unwrap().to_string(),
+		rivet_api_types::namespaces::runner_configs::RunnerConfig {
+			kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
+			metadata: None,
+			drain_on_version_upgrade: true,
+		},
 	);
-	// cursor is optional in pagination responses
+
+	crate::common::api::public::runner_configs_upsert(
+		dc.guard_port(),
+		rivet_api_peer::runner_configs::UpsertPath {
+			runner_name: envoy.pool_name().to_string(),
+		},
+		rivet_api_peer::runner_configs::UpsertQuery {
+			namespace: namespace.to_string(),
+		},
+		rivet_api_public::runner_configs::upsert::UpsertRequest { datacenters },
+	)
+	.await
+	.expect("failed to upsert runner config");
+
+	envoy.start().await.expect("failed to start envoy");
+	envoy.wait_ready().await;
+
+	envoy
 }
 
-// Datacenter helpers
-pub fn get_test_datacenter_names(_ctx: &super::TestCtx) -> Vec<String> {
-	vec!["dc-1".to_string(), "dc-2".to_string()] // Adjust based on actual DC names
+pub fn convert_strings_to_ids(actor_ids: Vec<String>) -> Vec<rivet_util::Id> {
+	actor_ids
+		.iter()
+		.map(|x| rivet_util::Id::from_str(&x).expect("failed to convert actor ids to string"))
+		.collect::<Vec<_>>()
 }
 
-pub async fn setup_multi_datacenter_test() -> super::TestCtx {
-	super::TestCtx::new_multi(2)
-		.await
-		.expect("Failed to setup multi-datacenter test")
+pub async fn create_actor(
+	port: u16,
+	namespace: &str,
+	name: &str,
+	runner_name: &str,
+	crash_policy: rivet_types::actors::CrashPolicy,
+) -> super::api_types::actors::create::CreateResponse {
+	super::api::public::actors_create(
+		port,
+		super::api_types::actors::create::CreateQuery {
+			namespace: namespace.to_string(),
+		},
+		super::api_types::actors::create::CreateRequest {
+			datacenter: None,
+			name: name.to_string(),
+			key: None,
+			input: None,
+			runner_name_selector: runner_name.to_string(),
+			crash_policy,
+		},
+	)
+	.await
+	.expect("failed to create actor")
+}
+
+pub fn generate_dummy_rivet_id(dc: &super::TestDatacenter) -> rivet_util::Id {
+	rivet_util::Id::new_v1(dc.config.dc_label())
 }

package/packages/engine/tests/common/test_runner.rs

@@ -0,0 +1,273 @@
+//! Test runner wrapper for engine tests.
+//!
+//! This module now adapts the Rust `rivet-test-envoy` harness to the legacy
+//! runner-oriented test surface that the engine tests still import.
+
+use anyhow::Result;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+pub use rivet_envoy_protocol::PROTOCOL_VERSION;
+pub use rivet_runner_protocol as protocol_types;
+pub use rivet_test_envoy::{
+	ActorConfig, ActorEvent, ActorLifecycleEvent, ActorStartResult, ActorStopResult,
+	CountingCrashActor, CrashNTimesThenSucceedActor, CrashOnStartActor, CustomActor,
+	CustomActorBuilder, DelayedStartActor, EchoActor, Envoy, EnvoyBuilder as TestEnvoyBuilderImpl,
+	EnvoyConfig, KvRequest, NotifyOnStartActor, SleepImmediatelyActor, StopImmediatelyActor,
+	TestActor, TimeoutActor, VerifyInputActor,
+};
+
+type ActorFactory = Arc<dyn Fn(ActorConfig) -> Box<dyn TestActor> + Send + Sync>;
+
+pub type TestRunner = Runner;
+pub type RunnerBuilderLegacy = RunnerBuilder;
+
+#[derive(Clone)]
+pub struct RunnerConfig {
+	endpoint: String,
+	token: String,
+	namespace: String,
+	runner_name: String,
+	runner_key: String,
+	version: u32,
+	total_slots: u32,
+}
+
+impl RunnerConfig {
+	pub fn builder() -> RunnerConfigBuilder {
+		RunnerConfigBuilder::default()
+	}
+}
+
+#[derive(Default)]
+pub struct RunnerConfigBuilder {
+	endpoint: Option<String>,
+	token: Option<String>,
+	namespace: Option<String>,
+	runner_name: Option<String>,
+	runner_key: Option<String>,
+	version: Option<u32>,
+	total_slots: Option<u32>,
+}
+
+impl RunnerConfigBuilder {
+	pub fn endpoint(mut self, endpoint: impl Into<String>) -> Self {
+		self.endpoint = Some(endpoint.into());
+		self
+	}
+
+	pub fn token(mut self, token: impl Into<String>) -> Self {
+		self.token = Some(token.into());
+		self
+	}
+
+	pub fn namespace(mut self, namespace: impl Into<String>) -> Self {
+		self.namespace = Some(namespace.into());
+		self
+	}
+
+	pub fn runner_name(mut self, runner_name: impl Into<String>) -> Self {
+		self.runner_name = Some(runner_name.into());
+		self
+	}
+
+	pub fn runner_key(mut self, runner_key: impl Into<String>) -> Self {
+		self.runner_key = Some(runner_key.into());
+		self
+	}
+
+	pub fn version(mut self, version: u32) -> Self {
+		self.version = Some(version);
+		self
+	}
+
+	pub fn total_slots(mut self, total_slots: u32) -> Self {
+		self.total_slots = Some(total_slots);
+		self
+	}
+
+	pub fn build(self) -> Result<RunnerConfig> {
+		Ok(RunnerConfig {
+			endpoint: self
+				.endpoint
+				.ok_or_else(|| anyhow::anyhow!("endpoint is required"))?,
+			token: self.token.unwrap_or_else(|| "dev".to_string()),
+			namespace: self
+				.namespace
+				.ok_or_else(|| anyhow::anyhow!("namespace is required"))?,
+			runner_name: self
+				.runner_name
+				.unwrap_or_else(|| "test-runner".to_string()),
+			runner_key: self
+				.runner_key
+				.unwrap_or_else(|| format!("key-{:012x}", rand::random::<u64>())),
+			version: self.version.unwrap_or(1),
+			total_slots: self.total_slots.unwrap_or(100),
+		})
+	}
+}
+
+pub struct RunnerBuilder {
+	config: RunnerConfig,
+	actor_factories: HashMap<String, ActorFactory>,
+}
+
+impl RunnerBuilder {
+	pub fn new(config: RunnerConfig) -> Self {
+		Self {
+			config,
+			actor_factories: HashMap::new(),
+		}
+	}
+
+	pub fn with_actor_behavior<F>(mut self, actor_name: &str, factory: F) -> Self
+	where
+		F: Fn(ActorConfig) -> Box<dyn TestActor> + Send + Sync + 'static,
+	{
+		self.actor_factories
+			.insert(actor_name.to_string(), Arc::new(factory));
+		self
+	}
+
+	pub fn build(self) -> Result<Runner> {
+		let envoy_config = EnvoyConfig::builder()
+			.endpoint(&self.config.endpoint)
+			.token(&self.config.token)
+			.namespace(&self.config.namespace)
+			.pool_name(&self.config.runner_name)
+			.version(self.config.version)
+			.metadata(serde_json::json!({
+				"runner_key": self.config.runner_key,
+				"total_slots": self.config.total_slots,
+			}))
+			.build()?;
+
+		let mut builder = TestEnvoyBuilderImpl::new(envoy_config);
+		for (name, factory) in self.actor_factories {
+			builder = builder.with_actor_behavior(&name, move |config| factory(config));
+		}
+
+		Ok(Runner {
+			runner_id: format!("runner-{}", uuid::Uuid::new_v4()),
+			runner_name: self.config.runner_name,
+			envoy: builder.build()?,
+		})
+	}
+}
+
+pub struct Runner {
+	pub runner_id: String,
+	runner_name: String,
+	envoy: Envoy,
+}
+
+impl Runner {
+	pub async fn start(&self) -> Result<()> {
+		self.envoy.start().await
+	}
+
+	pub async fn wait_ready(&self) -> String {
+		self.envoy.wait_ready().await;
+		self.runner_id.clone()
+	}
+
+	pub async fn has_actor(&self, actor_id: &str) -> bool {
+		self.envoy.has_actor(actor_id).await
+	}
+
+	pub async fn get_actor_ids(&self) -> Vec<String> {
+		self.envoy.get_actor_ids().await
+	}
+
+	pub fn name(&self) -> &str {
+		&self.runner_name
+	}
+
+	pub fn subscribe_lifecycle_events(
+		&self,
+	) -> tokio::sync::broadcast::Receiver<ActorLifecycleEvent> {
+		self.envoy.subscribe_lifecycle_events()
+	}
+
+	pub async fn shutdown(&self) {
+		self.envoy.shutdown().await;
+	}
+
+	pub async fn crash(&self) {
+		self.envoy.crash().await;
+	}
+}
+
+/// Test-specific runner builder that integrates with TestDatacenter.
+pub struct TestRunnerBuilder {
+	namespace: String,
+	runner_name: String,
+	runner_key: String,
+	version: u32,
+	total_slots: u32,
+	actor_factories: HashMap<String, ActorFactory>,
+}
+
+impl TestRunnerBuilder {
+	pub fn new(namespace: &str) -> Self {
+		Self {
+			namespace: namespace.to_string(),
+			runner_name: "test-runner".to_string(),
+			runner_key: format!("key-{:012x}", rand::random::<u64>()),
+			version: 1,
+			total_slots: 100,
+			actor_factories: HashMap::new(),
+		}
+	}
+
+	pub fn with_runner_name(mut self, name: &str) -> Self {
+		self.runner_name = name.to_string();
+		self
+	}
+
+	pub fn with_runner_key(mut self, key: &str) -> Self {
+		self.runner_key = key.to_string();
+		self
+	}
+
+	pub fn with_version(mut self, version: u32) -> Self {
+		self.version = version;
+		self
+	}
+
+	pub fn with_total_slots(mut self, total_slots: u32) -> Self {
+		self.total_slots = total_slots;
+		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,
+	{
+		self.actor_factories
+			.insert(actor_name.to_string(), Arc::new(factory));
+		self
+	}
+
+	pub async fn build(self, dc: &super::TestDatacenter) -> Result<Runner> {
+		let endpoint = format!("http://127.0.0.1:{}", dc.guard_port());
+		let token = "dev".to_string();
+
+		let config = RunnerConfig::builder()
+			.endpoint(&endpoint)
+			.token(&token)
+			.namespace(&self.namespace)
+			.runner_name(&self.runner_name)
+			.runner_key(&self.runner_key)
+			.version(self.version)
+			.total_slots(self.total_slots)
+			.build()?;
+
+		let mut builder = RunnerBuilder::new(config);
+		for (name, factory) in self.actor_factories {
+			builder = builder.with_actor_behavior(&name, move |config| factory(config));
+		}
+
+		builder.build()
+	}
+}