[TaskCenter] Replace TestCoreEnv with TestCoreEnv2

crates/admin/src/cluster_controller/scheduler.rs

@@ -581,7 +581,7 @@ mod tests {
         HashSet, PartitionProcessorPlacementHints, Scheduler,
     };
     use restate_core::network::{ForwardingHandler, Incoming, MessageCollectorMockConnector};
-    use restate_core::{Metadata, TestCoreEnv2, TestCoreEnvBuilder2};
+    use restate_core::{Metadata, TestCoreEnv, TestCoreEnvBuilder};
     use restate_types::cluster::cluster_state::{
         AliveNode, ClusterState, DeadNode, NodeState, PartitionProcessorStatus, RunMode,
     };
@@ -614,7 +614,7 @@ mod tests {
 
     #[test(restate_core::test)]
     async fn empty_leadership_changes_dont_modify_plan() -> googletest::Result<()> {
-        let test_env = TestCoreEnv2::create_with_single_node(0, 0).await;
+        let test_env = TestCoreEnv::create_with_single_node(0, 0).await;
         let metadata_store_client = test_env.metadata_store_client.clone();
         let networking = test_env.networking.clone();
 
@@ -690,7 +690,7 @@ mod tests {
         let (tx, control_recv) = mpsc::channel(100);
         let connector = MessageCollectorMockConnector::new(10, tx.clone());
 
-        let mut builder = TestCoreEnvBuilder2::with_transport_connector(connector);
+        let mut builder = TestCoreEnvBuilder::with_transport_connector(connector);
         builder.router_builder.add_raw_handler(
             TargetName::ControlProcessors,
             // network messages going to my node is also written to `tx`

crates/admin/src/cluster_controller/service.rs

@@ -482,9 +482,8 @@ mod tests {
     use restate_core::network::{
         FailingConnector, Incoming, MessageHandler, MockPeerConnection, NetworkServerBuilder,
     };
-    use restate_core::{
-        NoOpMessageHandler, TaskCenter, TaskKind, TestCoreEnv2, TestCoreEnvBuilder2,
-    };
+    use restate_core::test_env::NoOpMessageHandler;
+    use restate_core::{TaskCenter, TaskKind, TestCoreEnv, TestCoreEnvBuilder};
     use restate_types::cluster::cluster_state::PartitionProcessorStatus;
     use restate_types::config::{AdminOptions, Configuration};
     use restate_types::health::HealthStatus;
@@ -500,7 +499,7 @@ mod tests {
     #[test(restate_core::test)]
     async fn manual_log_trim() -> anyhow::Result<()> {
         const LOG_ID: LogId = LogId::new(0);
-        let mut builder = TestCoreEnvBuilder2::with_incoming_only_connector();
+        let mut builder = TestCoreEnvBuilder::with_incoming_only_connector();
         let bifrost_svc = BifrostService::new().with_factory(memory_loglet::Factory::default());
         let bifrost = bifrost_svc.handle();
 
@@ -775,11 +774,11 @@ mod tests {
     async fn create_test_env<F>(
         config: Configuration,
         mut modify_builder: F,
-    ) -> anyhow::Result<(TestCoreEnv2<FailingConnector>, Bifrost)>
+    ) -> anyhow::Result<(TestCoreEnv<FailingConnector>, Bifrost)>
     where
-        F: FnMut(TestCoreEnvBuilder2<FailingConnector>) -> TestCoreEnvBuilder2<FailingConnector>,
+        F: FnMut(TestCoreEnvBuilder<FailingConnector>) -> TestCoreEnvBuilder<FailingConnector>,
     {
-        let mut builder = TestCoreEnvBuilder2::with_incoming_only_connector();
+        let mut builder = TestCoreEnvBuilder::with_incoming_only_connector();
         let bifrost_svc = BifrostService::new().with_factory(memory_loglet::Factory::default());
         let bifrost = bifrost_svc.handle();
 

crates/bifrost/src/bifrost.rs

@@ -496,8 +496,8 @@ mod tests {
     use tracing::info;
     use tracing_test::traced_test;
 
-    use restate_core::TestCoreEnvBuilder2;
-    use restate_core::{TaskCenter, TaskKind, TestCoreEnv2};
+    use restate_core::TestCoreEnvBuilder;
+    use restate_core::{TaskCenter, TaskKind, TestCoreEnv};
     use restate_rocksdb::RocksDbManager;
     use restate_types::config::CommonOptions;
     use restate_types::live::Constant;
@@ -514,7 +514,7 @@ mod tests {
     #[traced_test]
     async fn test_append_smoke() -> googletest::Result<()> {
         let num_partitions = 5;
-        let _ = TestCoreEnvBuilder2::with_incoming_only_connector()
+        let _ = TestCoreEnvBuilder::with_incoming_only_connector()
             .set_partition_table(PartitionTable::with_equally_sized_partitions(
                 Version::MIN,
                 num_partitions,
@@ -586,7 +586,7 @@ mod tests {
 
     #[restate_core::test(start_paused = true)]
     async fn test_lazy_initialization() -> googletest::Result<()> {
-        let _ = TestCoreEnv2::create_with_single_node(1, 1).await;
+        let _ = TestCoreEnv::create_with_single_node(1, 1).await;
         let delay = Duration::from_secs(5);
         // This memory provider adds a delay to its loglet initialization, we want
         // to ensure that appends do not fail while waiting for the loglet;
@@ -604,7 +604,7 @@ mod tests {
     #[test(restate_core::test(flavor = "multi_thread", worker_threads = 2))]
     async fn trim_log_smoke_test() -> googletest::Result<()> {
         const LOG_ID: LogId = LogId::new(0);
-        let node_env = TestCoreEnvBuilder2::with_incoming_only_connector()
+        let node_env = TestCoreEnvBuilder::with_incoming_only_connector()
             .set_provider_kind(ProviderKind::Local)
             .build()
             .await;
@@ -676,7 +676,7 @@ mod tests {
     #[restate_core::test(start_paused = true)]
     async fn test_read_across_segments() -> googletest::Result<()> {
         const LOG_ID: LogId = LogId::new(0);
-        let node_env = TestCoreEnvBuilder2::with_incoming_only_connector()
+        let node_env = TestCoreEnvBuilder::with_incoming_only_connector()
             .set_partition_table(PartitionTable::with_equally_sized_partitions(
                 Version::MIN,
                 1,
@@ -862,7 +862,7 @@ mod tests {
     #[traced_test]
     async fn test_appends_correctly_handle_reconfiguration() -> googletest::Result<()> {
         const LOG_ID: LogId = LogId::new(0);
-        let node_env = TestCoreEnvBuilder2::with_incoming_only_connector()
+        let node_env = TestCoreEnvBuilder::with_incoming_only_connector()
             .set_partition_table(PartitionTable::with_equally_sized_partitions(
                 Version::MIN,
                 1,

crates/bifrost/src/providers/local_loglet/mod.rs

@@ -282,7 +282,7 @@ mod tests {
     use test_log::test;
 
     use crate::loglet::Loglet;
-    use restate_core::{TaskCenter, TestCoreEnvBuilder2};
+    use restate_core::{TaskCenter, TestCoreEnvBuilder};
     use restate_rocksdb::RocksDbManager;
     use restate_types::config::Configuration;
     use restate_types::live::Live;
@@ -307,7 +307,7 @@ mod tests {
     }
 
     async fn create_loglet() -> anyhow::Result<Arc<LocalLoglet>> {
-        let _node_env = TestCoreEnvBuilder2::with_incoming_only_connector()
+        let _node_env = TestCoreEnvBuilder::with_incoming_only_connector()
             .set_provider_kind(ProviderKind::Local)
             .build()
             .await;
@@ -345,7 +345,7 @@ mod tests {
 
     #[restate_core::test(flavor = "multi_thread", worker_threads = 4)]
     async fn local_loglet_append_after_seal_concurrent() -> googletest::Result<()> {
-        let _node_env = TestCoreEnvBuilder2::with_incoming_only_connector()
+        let _node_env = TestCoreEnvBuilder::with_incoming_only_connector()
             .set_provider_kind(ProviderKind::Local)
             .build()
             .await;

crates/bifrost/src/providers/memory_loglet.rs

@@ -403,14 +403,14 @@ impl Loglet for MemoryLoglet {
 mod tests {
     use super::*;
 
-    use restate_core::TestCoreEnvBuilder2;
+    use restate_core::TestCoreEnvBuilder;
 
     macro_rules! run_test {
         ($test:ident) => {
             paste::paste! {
                 #[restate_core::test(start_paused = true)]
                 async fn [<memory_loglet_  $test>]() -> googletest::Result<()> {
-                    let _node_env = TestCoreEnvBuilder2::with_incoming_only_connector()
+                    let _node_env = TestCoreEnvBuilder::with_incoming_only_connector()
                         .set_provider_kind(ProviderKind::InMemory)
                         .build()
                         .await;

crates/bifrost/src/providers/replicated_loglet/loglet.rs

@@ -345,7 +345,7 @@ mod tests {
     use test_log::test;
 
     use restate_core::network::NetworkServerBuilder;
-    use restate_core::{TaskCenter, TestCoreEnvBuilder2};
+    use restate_core::{TaskCenter, TestCoreEnvBuilder};
     use restate_log_server::LogServerService;
     use restate_rocksdb::RocksDbManager;
     use restate_types::config::{set_current_config, Configuration};
@@ -376,7 +376,7 @@ mod tests {
         let config = Live::from_value(config);
 
         let mut node_env =
-            TestCoreEnvBuilder2::with_incoming_only_connector().add_mock_nodes_config();
+            TestCoreEnvBuilder::with_incoming_only_connector().add_mock_nodes_config();
         let mut server_builder = NetworkServerBuilder::default();
 
         let logserver_rpc = LogServersRpc::new(&mut node_env.router_builder);

crates/bifrost/src/providers/replicated_loglet/remote_sequencer.rs

@@ -493,7 +493,7 @@ mod test {
 
     use restate_core::{
         network::{Incoming, MessageHandler, MockConnector},
-        TestCoreEnv2, TestCoreEnvBuilder2,
+        TestCoreEnv, TestCoreEnvBuilder,
     };
     use restate_types::{
         logs::{LogId, LogletOffset, Record, SequenceNumber, TailState},
@@ -553,7 +553,7 @@ mod test {
     }
 
     struct TestEnv {
-        pub core_env: TestCoreEnv2<MockConnector>,
+        pub core_env: TestCoreEnv<MockConnector>,
         pub remote_sequencer: RemoteSequencer<MockConnector>,
     }
 
@@ -565,7 +565,7 @@ mod test {
         let (connector, _receiver) = MockConnector::new(100);
         let connector = Arc::new(connector);
 
-        let mut builder = TestCoreEnvBuilder2::with_transport_connector(Arc::clone(&connector))
+        let mut builder = TestCoreEnvBuilder::with_transport_connector(Arc::clone(&connector))
             .add_mock_nodes_config()
             .add_message_handler(sequencer);
 

crates/bifrost/src/read_stream.rs

@@ -445,7 +445,7 @@ mod tests {
     use tracing::info;
     use tracing_test::traced_test;
 
-    use restate_core::{MetadataKind, TargetVersion, TaskCenter, TaskKind, TestCoreEnvBuilder2};
+    use restate_core::{MetadataKind, TargetVersion, TaskCenter, TaskKind, TestCoreEnvBuilder};
     use restate_rocksdb::RocksDbManager;
     use restate_types::config::{CommonOptions, Configuration};
     use restate_types::live::{Constant, Live};
@@ -461,7 +461,7 @@ mod tests {
     async fn test_readstream_one_loglet() -> anyhow::Result<()> {
         const LOG_ID: LogId = LogId::new(0);
 
-        let _ = TestCoreEnvBuilder2::with_incoming_only_connector()
+        let _ = TestCoreEnvBuilder::with_incoming_only_connector()
             .set_provider_kind(ProviderKind::Local)
             .build()
             .await;
@@ -541,7 +541,7 @@ mod tests {
     async fn test_read_stream_with_trim() -> anyhow::Result<()> {
         const LOG_ID: LogId = LogId::new(0);
 
-        let node_env = TestCoreEnvBuilder2::with_incoming_only_connector()
+        let node_env = TestCoreEnvBuilder::with_incoming_only_connector()
             .set_provider_kind(ProviderKind::Local)
             .build()
             .await;
@@ -634,7 +634,7 @@ mod tests {
     async fn test_readstream_simple_multi_loglet() -> anyhow::Result<()> {
         const LOG_ID: LogId = LogId::new(0);
 
-        let node_env = TestCoreEnvBuilder2::with_incoming_only_connector()
+        let node_env = TestCoreEnvBuilder::with_incoming_only_connector()
             .set_provider_kind(ProviderKind::Local)
             .build()
             .await;
@@ -790,7 +790,7 @@ mod tests {
     async fn test_readstream_sealed_multi_loglet() -> anyhow::Result<()> {
         const LOG_ID: LogId = LogId::new(0);
 
-        let node_env = TestCoreEnvBuilder2::with_incoming_only_connector()
+        let node_env = TestCoreEnvBuilder::with_incoming_only_connector()
             .set_provider_kind(ProviderKind::Local)
             .build()
             .await;
@@ -908,7 +908,7 @@ mod tests {
     async fn test_readstream_prefix_trimmed() -> anyhow::Result<()> {
         const LOG_ID: LogId = LogId::new(0);
 
-        let node_env = TestCoreEnvBuilder2::with_incoming_only_connector()
+        let node_env = TestCoreEnvBuilder::with_incoming_only_connector()
             .set_provider_kind(ProviderKind::Local)
             .build()
             .await;

crates/core/src/lib.rs

@@ -51,13 +51,7 @@ pub use task_center::{
 };
 
 #[cfg(any(test, feature = "test-util"))]
-mod test_env;
+pub mod test_env;
 
 #[cfg(any(test, feature = "test-util"))]
-mod test_env2;
-
-#[cfg(any(test, feature = "test-util"))]
-pub use test_env::{create_mock_nodes_config, NoOpMessageHandler, TestCoreEnv, TestCoreEnvBuilder};
-
-#[cfg(any(test, feature = "test-util"))]
-pub use test_env2::{TestCoreEnv2, TestCoreEnvBuilder2};
+pub use test_env::{TestCoreEnv, TestCoreEnvBuilder};

crates/core/src/network/connection_manager.rs

@@ -778,12 +778,12 @@ mod tests {
     use restate_types::Version;
 
     use crate::network::MockPeerConnection;
-    use crate::{self as restate_core, TestCoreEnv2, TestCoreEnvBuilder2};
+    use crate::{self as restate_core, TestCoreEnv, TestCoreEnvBuilder};
 
     // Test handshake with a client
     #[restate_core::test]
     async fn test_hello_welcome_handshake() -> Result<()> {
-        let _env = TestCoreEnv2::create_with_single_node(1, 1).await;
+        let _env = TestCoreEnv::create_with_single_node(1, 1).await;
         let metadata = Metadata::current();
         let connections = ConnectionManager::new_incoming_only(metadata.clone());
 
@@ -802,7 +802,7 @@ mod tests {
 
     #[restate_core::test(start_paused = true)]
     async fn test_hello_welcome_timeout() -> Result<()> {
-        let _env = TestCoreEnv2::create_with_single_node(1, 1).await;
+        let _env = TestCoreEnv::create_with_single_node(1, 1).await;
         let metadata = Metadata::current();
         let net_opts = NetworkingOptions::default();
         let (_tx, rx) = mpsc::channel(1);
@@ -824,7 +824,7 @@ mod tests {
 
     #[restate_core::test]
     async fn test_bad_handshake() -> Result<()> {
-        let test_setup = TestCoreEnv2::create_with_single_node(1, 1).await;
+        let test_setup = TestCoreEnv::create_with_single_node(1, 1).await;
         let metadata = test_setup.metadata;
         let (tx, rx) = mpsc::channel(1);
         let my_node_id = metadata.my_node_id();
@@ -900,7 +900,7 @@ mod tests {
 
     #[restate_core::test]
     async fn test_node_generation() -> Result<()> {
-        let _env = TestCoreEnv2::create_with_single_node(1, 2).await;
+        let _env = TestCoreEnv::create_with_single_node(1, 2).await;
         let metadata = Metadata::current();
         let (tx, rx) = mpsc::channel(1);
         let mut my_node_id = metadata.my_node_id();
@@ -995,7 +995,7 @@ mod tests {
         );
         nodes_config.upsert_node(node_config);
 
-        let test_env = TestCoreEnvBuilder2::with_incoming_only_connector()
+        let test_env = TestCoreEnvBuilder::with_incoming_only_connector()
             .set_nodes_config(nodes_config)
             .build()
             .await;