From 9a7b061634a902fe94da77741dd30ab618fad352 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicol=C3=A1s=20Venturo?= <nicolas.venturo@gmail.com>
Date: Wed, 17 Apr 2024 17:12:28 -0300
Subject: [PATCH] test: use new mock.get_last_params() for public storage
 writes (#5823)

This uses the new mock features implemented in
https://github.com/noir-lang/noir/pull/4789 to test calls to storage
writes. It's not incredibly exciting, but is a nice example of how such
an oracle is used and nicely adds to the completeness of the
~`SharedMutable`~`PriblicMutable` tests.
---
 aztec/src/public_storage.nr                   | 10 ++--
 .../shared_mutable/shared_mutable.nr          | 52 ++++++++++++++++---
 2 files changed, 52 insertions(+), 10 deletions(-)

diff --git a/aztec/src/public_storage.nr b/aztec/src/public_storage.nr
index 7375ab91..d46b2c5f 100644
--- a/aztec/src/public_storage.nr
+++ b/aztec/src/public_storage.nr
@@ -61,8 +61,12 @@ mod tests {
 
     #[test]
     fn test_write() {
-        // Here we'd want to test that what is written to storage is deserialized to the same struct, but the current
-        // oracle mocks lack these capabilities.
-        // TODO: implement this once https://github.com/noir-lang/noir/issues/4652 is closed
+        let slot = 7;
+        let to_write = TestStruct { a: 13, b: 42 };
+
+        let mock = OracleMock::mock("storageWrite").returns([0; 2]); // The return value is unused
+
+        public_storage::write(slot, to_write);
+        assert_eq(mock.get_last_params(), (slot, to_write.serialize()));
     }
 }
diff --git a/aztec/src/state_vars/shared_mutable/shared_mutable.nr b/aztec/src/state_vars/shared_mutable/shared_mutable.nr
index 63e3a9c0..5d327b1f 100644
--- a/aztec/src/state_vars/shared_mutable/shared_mutable.nr
+++ b/aztec/src/state_vars/shared_mutable/shared_mutable.nr
@@ -155,7 +155,7 @@ mod test {
 
     #[test]
     fn test_get_current_value_in_public_after_change() {
-        let (state_var, block_number ) = setup(false);
+        let (state_var, block_number) = setup(false);
 
         let slot = state_var.get_derived_storage_slot();
         let (pre, post) = (13, 17);
@@ -172,16 +172,54 @@ mod test {
         let slot = state_var.get_derived_storage_slot();
         let (pre, post) = (13, 17);
 
+        // Change in the future
+        OracleMock::mock("storageRead").with_params((slot, 3)).returns([pre, post, block_number + 1]);
+
+        let write_mock = OracleMock::mock("storageWrite").returns([0; 3]); // The oracle return value is actually unused
+
+        let new_value = 42;
+        state_var.schedule_value_change(new_value);
+
+        // The new scheduled change replaces the old one
+        assert_eq(write_mock.get_last_params(), (slot, [pre, new_value, block_number + TEST_DELAY]));
+    }
+
+    #[test]
+    fn test_schedule_value_change_at_change() {
+        let (state_var, block_number) = setup(false);
+
         let slot = state_var.get_derived_storage_slot();
+        let (pre, post) = (13, 17);
 
-        OracleMock::mock("storageRead").with_params((slot, 3)).returns([pre, post, block_number + 1]);
+        // Change in the current block
+        OracleMock::mock("storageRead").with_params((slot, 3)).returns([pre, post, block_number]);
+
+        let write_mock = OracleMock::mock("storageWrite").returns([0; 3]); // The oracle return value is actually unused
+
+        let new_value = 42;
+        state_var.schedule_value_change(new_value);
+
+        // The previous 'post' value is the current one and becomes the 'pre' value
+        assert_eq(write_mock.get_last_params(), (slot, [post, new_value, block_number + TEST_DELAY]));
+    }
+
+    #[test]
+    fn test_schedule_value_change_after_change() {
+        let (state_var, block_number) = setup(false);
+
+        let slot = state_var.get_derived_storage_slot();
+        let (pre, post) = (13, 17);
+
+        // Change in the past
+        OracleMock::mock("storageRead").with_params((slot, 3)).returns([pre, post, block_number - 1]);
+
+        let write_mock = OracleMock::mock("storageWrite").returns([0; 3]); // The oracle return value is actually unused
 
         let new_value = 42;
-        // Here we want to assert that the `storageWrite` oracle is called with a certain set of values, but the current
-        // oracle mocks don't have those capabilities.
-        // TODO: implement this once https://github.com/noir-lang/noir/issues/4652 is closed
-        // OracleMock::mock("storageWrite").expect_call((slot, [pre, new_value, block_number + DELAY]));
-        // state_var.schedule_value_change(new_value);
+        state_var.schedule_value_change(new_value);
+
+        // The previous 'post' value is the current one and becomes the 'pre' value
+        assert_eq(write_mock.get_last_params(), (slot, [post, new_value, block_number + TEST_DELAY]));
     }
 
     #[test]