From b4ba793963ca92e9edf14ab1741b86bc70b0f61d Mon Sep 17 00:00:00 2001
From: Yan Chen <yan.chen@dfinity.org>
Date: Fri, 22 Sep 2023 21:13:57 -0700
Subject: [PATCH] test upgrade

---
 src/instrumentation.rs      |  2 +-
 tests/deployable.ic-repl.sh | 18 +++++++++++++++++-
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/instrumentation.rs b/src/instrumentation.rs
index 0f8ba44..4470a3e 100644
--- a/src/instrumentation.rs
+++ b/src/instrumentation.rs
@@ -613,7 +613,7 @@ fn inject_pre_upgrade(m: &mut Module, vars: &Variables, config: &Config) {
     let mut builder = get_or_create_export_func(m, "canister_pre_upgrade");
     #[rustfmt::skip]
     builder
-        // backup memory
+        // backup memory. This is not restricted needed, since the end of pre-upgrade.
         .i32_const(0)
         .load(memory, LoadKind::I64 { atomic: false }, MemArg { offset: 0, align: 8})
         .local_set(a)
diff --git a/tests/deployable.ic-repl.sh b/tests/deployable.ic-repl.sh
index 3a63452..d57cfbd 100644
--- a/tests/deployable.ic-repl.sh
+++ b/tests/deployable.ic-repl.sh
@@ -14,6 +14,16 @@ function install(wasm) {
   );
   S
 };
+function upgrade(S, wasm) {
+  call ic.install_code(
+    record {
+      arg = encode ();
+      wasm_module = wasm;
+      mode = variant { upgrade };
+      canister_id = S;
+    }
+  );
+};
 
 function motoko(wasm) {
   let S = install(wasm);
@@ -94,14 +104,20 @@ function check_profiling(S, cycles, len) {
   assert _ == (cycles : int64);
   call S.__get_profiling();
   assert _.size() == (len : nat);
+  null
 };
 
 let S = motoko(file("ok/motoko-instrument.wasm"));
 check_profiling(S, 9397, 78);
 let S = motoko(file("ok/motoko-gc-instrument.wasm"));
 check_profiling(S, 250, 4);
-let S = motoko(file("ok/motoko-region-instrument.wasm"));
+let wasm = file("ok/motoko-region-instrument.wasm");
+let S = motoko(wasm);
 check_profiling(S, 463589, 78);
+upgrade(S, wasm);
+call S.get();
+assert _ == (45 : nat);
+check_profiling(S, 472223, 374);
 motoko(file("ok/motoko-shrink.wasm"));
 motoko(file("ok/motoko-limit.wasm"));