diff --git a/rs/nns/integration_tests/src/canister_upgrade.rs b/rs/nns/integration_tests/src/canister_upgrade.rs
index 82b8a705beec..785a8c572760 100644
--- a/rs/nns/integration_tests/src/canister_upgrade.rs
+++ b/rs/nns/integration_tests/src/canister_upgrade.rs
@@ -45,6 +45,7 @@ fn test_upgrade_canister(canister_id: CanisterId, canister_wasm: Wasm) {
         canister_id,
         &Sha256::hash(&modified_wasm),
         controller_canister_id,
+        false,
     );
 
     let proposal_info =
diff --git a/rs/nns/integration_tests/src/upgrade_canisters_with_golden_nns_state.rs b/rs/nns/integration_tests/src/upgrade_canisters_with_golden_nns_state.rs
index 01cee0ad474d..903eaec5726b 100644
--- a/rs/nns/integration_tests/src/upgrade_canisters_with_golden_nns_state.rs
+++ b/rs/nns/integration_tests/src/upgrade_canisters_with_golden_nns_state.rs
@@ -255,6 +255,8 @@ fn test_upgrade_canisters_with_golden_nns_state() {
     // TODO: Use PocketIc instead of StateMachine.
     let state_machine = new_state_machine_with_golden_nns_state_or_panic();
 
+    state_machine.reject_remote_callbacks();
+
     // Step 1.2: Create a super powerful Neuron.
     println!("Creating super powerful Neuron.");
     let neuron_controller = PrincipalId::new_self_authenticating(&[1, 2, 3, 4]);
@@ -334,6 +336,7 @@ fn test_upgrade_canisters_with_golden_nns_state() {
                     *canister_id,
                     wasm_hash,
                     nns_canister_upgrade.controller_principal_id(),
+                    true,
                 );
                 println!(
                     "Attempt {repetition_number} to upgrade {nns_canister_name} was successful."
diff --git a/rs/nns/test_utils/src/state_test_helpers.rs b/rs/nns/test_utils/src/state_test_helpers.rs
index f9a6444a29ec..ec82fe3dea7c 100644
--- a/rs/nns/test_utils/src/state_test_helpers.rs
+++ b/rs/nns/test_utils/src/state_test_helpers.rs
@@ -1251,10 +1251,14 @@ pub fn wait_for_canister_upgrade_to_succeed(
     new_wasm_hash: &[u8; 32],
     // For most NNS canisters, ROOT_CANISTER_ID would be passed here (modulo conversion).
     controller_principal_id: PrincipalId,
+    reject_remote_callbacks_while_waiting: bool,
 ) {
     let mut last_status = None;
     for i in 0..25 {
         state_machine.tick();
+        if reject_remote_callbacks_while_waiting {
+            state_machine.reject_remote_callbacks();
+        }
 
         // Fetch status of the canister being upgraded.
         let status_result = get_canister_status(