fix tests

Signed-off-by: Bugen Zhao <[email protected]>

Author: BugenZhao

src/tests/simulation/src/cluster.rs

@@ -472,8 +472,8 @@ impl Cluster {
             .await
     }
 
-    /// Kill some nodes and restart them in 2s + restart_delay_secs with a probability of 0.1.
-    #[cfg_or_panic(madsim)]
+    /// Generate a list of random worker nodes to kill by `opts`, then call `kill_nodes` to kill and
+    /// restart them.
     pub async fn kill_node(&self, opts: &KillOpts) {
         let mut nodes = vec![];
         if opts.kill_meta {
@@ -528,7 +528,20 @@ impl Cluster {
                 nodes.push(format!("compactor-{}", i));
             }
         }
-        join_all(nodes.iter().map(|name| async move {
+
+        self.kill_nodes(nodes, opts.restart_delay_secs).await
+    }
+
+    /// Kill the given nodes by their names and restart them in 2s + restart_delay_secs with a
+    /// probability of 0.1.
+    #[cfg_or_panic(madsim)]
+    pub async fn kill_nodes(
+        &self,
+        nodes: impl IntoIterator<Item = impl AsRef<str>>,
+        restart_delay_secs: u32,
+    ) {
+        join_all(nodes.into_iter().map(|name| async move {
+            let name = name.as_ref();
             let t = rand::thread_rng().gen_range(Duration::from_secs(0)..Duration::from_secs(1));
             tokio::time::sleep(t).await;
             tracing::info!("kill {name}");
@@ -540,7 +553,7 @@ impl Cluster {
             // so that the node is expired and removed from the cluster
             if rand::thread_rng().gen_bool(0.1) {
                 // max_heartbeat_interval_secs = 15
-                t += Duration::from_secs(opts.restart_delay_secs as u64);
+                t += Duration::from_secs(restart_delay_secs as u64);
             }
             tokio::time::sleep(t).await;
             tracing::info!("restart {name}");

src/tests/simulation/tests/integration_tests/recovery/pause_on_bootstrap.rs

@@ -1,7 +1,21 @@
+// Copyright 2023 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 use std::time::Duration;
 
 use anyhow::Result;
-use risingwave_simulation::cluster::{Configuration, KillOpts};
+use risingwave_simulation::cluster::Configuration;
 use risingwave_simulation::nexmark::NexmarkCluster;
 use risingwave_simulation::utils::AssertResult;
 use tokio::time::{sleep, timeout};
@@ -24,7 +38,16 @@ enum ResumeBy {
 }
 
 async fn test_impl(resume_by: ResumeBy) -> Result<()> {
-    let mut cluster = NexmarkCluster::new(Configuration::for_scale(), 6, None, false).await?;
+    let mut cluster = NexmarkCluster::new(
+        Configuration {
+            meta_nodes: 1,
+            ..Configuration::for_scale()
+        },
+        6,
+        None,
+        false,
+    )
+    .await?;
 
     cluster.run(SET_PARAMETER).await?;
     cluster.run(CREATE).await?;
@@ -33,8 +56,8 @@ async fn test_impl(resume_by: ResumeBy) -> Result<()> {
     // Run for a while.
     sleep(Duration::from_secs(10)).await;
 
-    // Kill all nodes and wait for the service to recover.
-    cluster.kill_node(&KillOpts::ALL).await;
+    // Kill the meta node and wait for the service to recover.
+    cluster.kill_nodes(["meta-1"], 0).await;
     sleep(Duration::from_secs(10)).await;
 
     // The source should be paused.
@@ -64,20 +87,28 @@ async fn test_impl(resume_by: ResumeBy) -> Result<()> {
 
     match resume_by {
         ResumeBy::Risectl => cluster.resume().await?,
-        ResumeBy::Restart => cluster.kill_node(&KillOpts::ALL).await,
+        ResumeBy::Restart => cluster.kill_nodes(["meta-1"], 0).await,
     }
     sleep(Duration::from_secs(10)).await;
 
     // The source should be resumed.
     let new_count = cluster.run(SELECT).await?;
     assert_ne!(count, new_count);
 
-    // DML on tables should be allowed.
+    // DML on tables should be allowed. However, we're uncertain whether the previous blocked DML is
+    // executed or not. So we just check the count difference.
     {
         let mut session = cluster.start_session();
+
+        session.run("FLUSH").await?;
+        let count: i64 = session.run(SELECT_COUNT_TABLE).await?.parse().unwrap();
+
         session.run(INSERT_INTO_TABLE).await?;
         session.run("FLUSH").await?;
-        session.run(SELECT_COUNT_TABLE).await?.assert_result_eq("1");
+        session
+            .run(SELECT_COUNT_TABLE)
+            .await?
+            .assert_result_eq(format!("{}", count + 1));
     }
 
     Ok(())