diff --git a/rs/nns/handlers/root/impl/test_canisters/upgrade_test_canister.rs b/rs/nns/handlers/root/impl/test_canisters/upgrade_test_canister.rs
index d6966caae8ec..021a2840a88d 100644
--- a/rs/nns/handlers/root/impl/test_canisters/upgrade_test_canister.rs
+++ b/rs/nns/handlers/root/impl/test_canisters/upgrade_test_canister.rs
@@ -3,10 +3,15 @@
 use ic_cdk::{
     api::{
         call::arg_data_raw,
-        stable::{stable_bytes, stable_write},
+        stable::{stable_grow, stable_read, stable_write},
     },
     post_upgrade, println, query,
 };
+use std::cell::RefCell;
+
+thread_local! {
+    static ARG_LEN: RefCell<usize>  = RefCell::new(0);
+}
 
 fn main() {}
 
@@ -14,10 +19,18 @@ fn main() {}
 fn post_upgrade() {
     let arg = arg_data_raw();
     println!("Initializing test canister with arg={:?}", arg);
+    stable_grow(1).expect("Could not grow stable memory");
+    ARG_LEN.with(|len| {
+        *len.borrow_mut() = arg.len();
+    });
     stable_write(0, &arg);
 }
 
 #[query]
 fn read_stable() -> Vec<u8> {
-    stable_bytes()
+    let len = ARG_LEN.with(|len| *len.borrow());
+    let mut buf = vec![0; len];
+    stable_read(0, &mut buf);
+
+    buf
 }
diff --git a/rs/nns/handlers/root/impl/tests/test_reinstall_and_upgrade.rs b/rs/nns/handlers/root/impl/tests/test_reinstall_and_upgrade.rs
index b437384c87bd..135556d65648 100644
--- a/rs/nns/handlers/root/impl/tests/test_reinstall_and_upgrade.rs
+++ b/rs/nns/handlers/root/impl/tests/test_reinstall_and_upgrade.rs
@@ -1,7 +1,7 @@
 use assert_matches::assert_matches;
 use candid::Encode;
 use canister_test::{Project, Runtime};
-use dfn_candid::candid;
+use dfn_candid::{candid, candid_one};
 use ic_management_canister_types::CanisterInstallMode::{self, Install, Reinstall, Upgrade};
 use ic_nervous_system_clients::{
     canister_id_record::CanisterIdRecord,
@@ -192,7 +192,7 @@ fn test_init_payload_is_passed_through_upgrades() {
         );
 
         // Now let's wait for the upgrade to complete
-        loop {
+        for _ in 1..100 {
             let status: CanisterStatusResult = root
                 .update_(
                     "canister_status",
@@ -206,13 +206,11 @@ fn test_init_payload_is_passed_through_upgrades() {
             }
         }
 
-        assert_eq!(
-            universal
-                .query_("read_stable", bytes, Vec::new())
-                .await
-                .unwrap(),
-            test_byte_array,
-        );
+        let response_bytes: Vec<u8> = universal
+            .query_("read_stable", candid_one::<Vec<u8>, ()>, ())
+            .await
+            .unwrap();
+        assert_eq!(response_bytes, test_byte_array,);
 
         Ok(())
     });