Split SBM and MBM check opts

README.md

@@ -136,9 +136,11 @@ after the migration. [`OnRuntimeUpgrade::pre_upgrade`] returns a [`Vec<u8>`] tha
 arbitrary encoded data (usually some pre-upgrade state) which will be passed to
 [`OnRuntimeUpgrade::pre_upgrade`] after upgrading and used for post checking.
 
-**Note on Multi-Block Migrations (MBM):** If the runtime uses MBMs, the standard
-`pre_upgrade` and `post_upgrade` checks might be skipped by the executive. To
-force these hooks to run synchronously for testing, use the `--disable-mbm-checks` flag.
+`--checks` (default `pre-and-post`) selects the checks run after multi-block migrations (MBMs),
+or around single-block migrations (SBMs) with `--disable-mbm-checks`. To select both phases
+separately, use `--sbm-checks` and `--mbm-checks` instead, e.g.
+`--sbm-checks=pre-and-post --mbm-checks=try-state`. `--mbm-checks=pre-and-post` runs the SBM
+hooks again on the migrated state. The MBMs' own pre/post hooks always run.
 
 ### [`VersionedMigration`]
 

cli/main.rs

@@ -152,9 +152,11 @@
 //! arbitrary encoded data (usually some pre-upgrade state) which will be passed to
 //! [`OnRuntimeUpgrade::pre_upgrade`] after upgrading and used for post checking.
 //!
-//! **Note on Multi-Block Migrations (MBM):** If the runtime uses MBMs, the standard
-//! `pre_upgrade` and `post_upgrade` checks might be skipped by the executive. To
-//! force these hooks to run synchronously for testing, use the `--disable-mbm-checks` flag.
+//! `--checks` (default `pre-and-post`) selects the checks run after multi-block migrations (MBMs),
+//! or around single-block migrations (SBMs) with `--disable-mbm-checks`. To select both phases
+//! separately, use `--sbm-checks` and `--mbm-checks` instead, e.g.
+//! `--sbm-checks=pre-and-post --mbm-checks=try-state`. `--mbm-checks=pre-and-post` runs the SBM
+//! hooks again on the migrated state. The MBMs' own pre/post hooks always run.
 //!
 //! ### [`VersionedMigration`]
 //!

cli/tests/on_runtime_upgrade.rs

@@ -169,6 +169,36 @@ mod on_runtime_upgrade {
         .await;
     }
 
+    #[tokio::test]
+    async fn pre_upgrade_fail_sbm_checks_fails() {
+        run_test(
+            TestConfig::new("rococo-people", "people_rococo_runtime_pre_upgrade_fail")
+                .with_sub_command_args(&["--sbm-checks=pre-and-post", "--mbm-checks=none"]),
+            false,
+        )
+        .await;
+    }
+
+    #[tokio::test]
+    async fn pre_upgrade_fail_mbm_checks_fails() {
+        run_test(
+            TestConfig::new("rococo-people", "people_rococo_runtime_pre_upgrade_fail")
+                .with_sub_command_args(&["--sbm-checks=none", "--mbm-checks=pre-and-post"]),
+            false,
+        )
+        .await;
+    }
+
+    #[tokio::test]
+    async fn pre_upgrade_fail_sbm_and_mbm_checks_can_be_ignored() {
+        run_test(
+            TestConfig::new("rococo-people", "people_rococo_runtime_pre_upgrade_fail")
+                .with_sub_command_args(&["--sbm-checks=none", "--mbm-checks=none"]),
+            true,
+        )
+        .await;
+    }
+
     #[tokio::test]
     async fn post_upgrade_fail_fails() {
         run_test(
@@ -383,8 +413,8 @@ mod on_runtime_upgrade {
     async fn mbm_fail_ignore_mbms() {
         run_test(
             TestConfig::new("rococo-people", "people_rococo_runtime_mbm_fails")
-                .with_sub_command_args(&["--disable-mbms-checks"]),
-            false,
+                .with_sub_command_args(&["--disable-mbm-checks"]),
+            true,
         )
         .await;
     }

core/src/commands/on_runtime_upgrade/mbms.rs

@@ -47,7 +47,7 @@ where
             Level::Info,
             &format!(
                 "🔬 Running Multi-Block-Migrations with checks: {:?}",
-                self.command.checks
+                self.command.multi_block_checks()
             ),
         );
 
@@ -59,7 +59,7 @@ where
             .await?;
 
         if core_version::<Block, HostFns>(&ext, &executor)? < 5 {
-            return Err("Your runtime does not support Multi-Block-Migrations. Please disable the check with `--mbms false` or update your runtime.".into());
+            return Err("Your runtime does not support Multi-Block-Migrations. Please disable the check with `--disable-mbm-checks` or update your runtime.".into());
         }
 
         let inner_ext = Arc::new(Mutex::new(ext.inner_ext));
@@ -114,7 +114,7 @@ where
             &mut Default::default(),
             &executor,
             "TryRuntime_on_runtime_upgrade",
-            self.command.checks.encode().as_ref(),
+            self.command.multi_block_checks().encode().as_ref(),
             Default::default(), // TODO
             None,
         )?;

core/src/commands/on_runtime_upgrade/mod.rs

@@ -60,16 +60,29 @@ pub struct Command {
     ///   present).
     /// - `try-state`: Perform the try-state checks.
     ///
-    /// Performing any checks will potentially invalidate the measured PoV/Weight.
+    /// Runs after MBMs, or around SBMs with --disable-mbm-checks. Performing any checks will
+    /// potentially invalidate the measured PoV/Weight.
     // NOTE: The clap attributes make it backwards compatible with the previous `--checks` flag.
     #[clap(long,
 		default_value = "pre-and-post",
 		default_missing_value = "all",
 		num_args = 0..=1,
+		conflicts_with_all = ["sbm_checks", "mbm_checks"],
 		verbatim_doc_comment
     )]
     pub checks: UpgradeCheckSelect,
 
+    /// Checks around single-block migrations, like --checks. Requires --mbm-checks unless
+    /// --disable-mbm-checks is set.
+    #[clap(long, default_missing_value = "all", num_args = 0..=1, requires = "mbm_checks")]
+    pub sbm_checks: Option<UpgradeCheckSelect>,
+
+    /// Checks after multi-block migrations finish, like --checks. Does not affect the MBMs' own
+    /// pre/post hooks. Requires --sbm-checks.
+    #[clap(long, default_missing_value = "all", num_args = 0..=1, requires = "sbm_checks",
+        conflicts_with = "disable_mbm_checks")]
+    pub mbm_checks: Option<UpgradeCheckSelect>,
+
     /// Whether to disable weight warnings. Useful if a relay chain's runtime is being tested.
     #[clap(long, default_value = "false", default_missing_value = "true")]
     pub no_weight_warnings: bool,
@@ -104,6 +117,84 @@ pub struct Command {
     pub blocktime: u64,
 }
 
+impl Command {
+    fn single_block_checks(&self) -> UpgradeCheckSelect {
+        match self.sbm_checks {
+            Some(checks) => checks,
+            None if self.disable_mbm_checks => self.checks,
+            None => UpgradeCheckSelect::None,
+        }
+    }
+
+    fn multi_block_checks(&self) -> UpgradeCheckSelect {
+        match (self.sbm_checks, self.mbm_checks) {
+            (_, Some(checks)) => checks,
+            (Some(_), None) => UpgradeCheckSelect::None,
+            (None, None) => self.checks,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use clap::{error::ErrorKind, Parser};
+    use UpgradeCheckSelect::{All, PreAndPost, TryState};
+
+    use super::*;
+
+    fn parse(args: &[&str]) -> Result<Command, clap::Error> {
+        Command::try_parse_from(
+            ["on-runtime-upgrade"]
+                .into_iter()
+                .chain(args.iter().copied())
+                .chain(["--blocktime=6000", "snap", "--path=unused.snap"]),
+        )
+    }
+
+    #[test]
+    fn check_selection() {
+        for (args, sbm, mbm) in [
+            (&[][..], UpgradeCheckSelect::None, PreAndPost),
+            (&["--checks"], UpgradeCheckSelect::None, All),
+            (&["--checks=all", "--disable-mbm-checks"], All, All),
+            (
+                &["--sbm-checks=pre-and-post", "--mbm-checks=try-state"],
+                PreAndPost,
+                TryState,
+            ),
+            (
+                &["--sbm-checks", "--disable-mbm-checks"],
+                All,
+                UpgradeCheckSelect::None,
+            ),
+        ] {
+            let command = parse(args).unwrap();
+            assert_eq!(
+                (command.single_block_checks(), command.multi_block_checks()),
+                (sbm, mbm)
+            );
+        }
+    }
+
+    #[test]
+    fn invalid_check_selection() {
+        for (args, kind) in [
+            (&["--sbm-checks"][..], ErrorKind::MissingRequiredArgument),
+            (&["--mbm-checks"], ErrorKind::MissingRequiredArgument),
+            (
+                &["--checks", "--sbm-checks", "--mbm-checks"],
+                ErrorKind::ArgumentConflict,
+            ),
+            (
+                &["--sbm-checks", "--mbm-checks", "--disable-mbm-checks"],
+                ErrorKind::ArgumentConflict,
+            ),
+        ] {
+            assert_eq!(parse(args).unwrap_err().kind(), kind, "{args:?}");
+        }
+    }
+}
+
 /// Convenience struct to hold all the generic args and where clauses.
 pub(crate) struct CheckOnRuntimeUpgrade<Block, HostFns> {
     pub shared: SharedParams,
@@ -136,11 +227,7 @@ where
             .to_ext::<Block, HostFns>(shared, &executor, None, runtime_checks)
             .await?;
 
-        let sync_checks = if command.disable_mbm_checks {
-            command.checks
-        } else {
-            UpgradeCheckSelect::None
-        };
+        let sync_checks = command.single_block_checks();
 
         // Run `TryRuntime_on_runtime_upgrade` with the given checks.
         basti_log(
@@ -187,11 +274,7 @@ where
         if !self.command.disable_idempotency_checks {
             basti_log(
                 Level::Info,
-                format!(
-                    "🔬 Running TryRuntime_on_runtime_upgrade again to check idempotency: {:?}",
-                    self.command.checks
-                )
-                .as_str(),
+                "🔬 Running TryRuntime_on_runtime_upgrade again to check idempotency",
             );
             let executor = build_executor(&self.shared);