feat: allow patch state to take in multiple patches + access key and account patching

examples/src/spooning.rs

@@ -97,11 +97,7 @@ async fn main() -> anyhow::Result<()> {
 
     // Patch our testnet STATE into our local sandbox:
     worker
-        .patch_state(
-            sandbox_contract.id(),
-            "STATE".as_bytes(),
-            &status_msg.try_to_vec()?,
-        )
+        .patch_state(sandbox_contract.id(), "STATE", status_msg.try_to_vec()?)
         .await?;
 
     // Now grab the state to see that it has indeed been patched:

workspaces/build.rs

@@ -1,8 +1,9 @@
 fn main() {
     let doc_build = cfg!(doc) || std::env::var("DOCS_RS").is_ok();
     if !doc_build && cfg!(feature = "install") {
-        // TODO Update commit to stable version once binaries are published correctly
-        near_sandbox_utils::install_with_version("master/97c0410de519ecaca369aaee26f0ca5eb9e7de06")
-            .expect("Could not install sandbox");
+        match near_sandbox_utils::ensure_sandbox_bin(){
+            Ok(p) => println!("Successfully installed sandbox in: {:?}", p),
+            Err(e) => panic!("Could not install sandbox\nReason: {:?}", e),
+        }
     }
 }

workspaces/src/network/mod.rs

@@ -10,6 +10,9 @@ mod server;
 mod testnet;
 pub(crate) mod variants;
 
+pub(crate) use sandbox::PatchAccessKeyTransaction;
+pub(crate) use sandbox::PatchStateAccountTransaction;
+pub(crate) use sandbox::PatchStateTransaction;
 pub(crate) use variants::DEV_ACCOUNT_SEED;
 
 pub use self::betanet::Betanet;

workspaces/src/network/sandbox.rs

@@ -4,15 +4,18 @@ use std::str::FromStr;
 use async_trait::async_trait;
 use near_jsonrpc_client::methods::sandbox_fast_forward::RpcSandboxFastForwardRequest;
 use near_jsonrpc_client::methods::sandbox_patch_state::RpcSandboxPatchStateRequest;
+use near_primitives::hash::CryptoHash;
 use near_primitives::state_record::StateRecord;
+use near_primitives::views::AccountView;
+use std::iter::IntoIterator;
 
 use super::{AllowDevAccountCreation, NetworkClient, NetworkInfo, TopLevelAccountCreator};
 use crate::network::server::SandboxServer;
 use crate::network::Info;
 use crate::result::CallExecution;
 use crate::rpc::client::Client;
 use crate::rpc::patch::ImportContractTransaction;
-use crate::types::{AccountId, Balance, InMemorySigner, SecretKey};
+use crate::types::{AccountId, Balance, InMemorySigner, Nonce, SecretKey, StorageUsage};
 use crate::{Account, Contract, Network, Worker};
 
 // Constant taken from nearcore crate to avoid dependency
@@ -135,36 +138,258 @@ impl Sandbox {
         ImportContractTransaction::new(id.to_owned(), worker.client(), self.client())
     }
 
-    pub(crate) async fn patch_state(
+    pub(crate) fn patch_state(&self, account_id: AccountId) -> PatchStateTransaction {
+        PatchStateTransaction::new(self, account_id)
+    }
+
+    pub(crate) fn patch_account(&self, account_id: AccountId) -> PatchStateAccountTransaction {
+        PatchStateAccountTransaction::new(self, account_id)
+    }
+
+    pub(crate) fn patch_access_key(
         &self,
-        contract_id: &AccountId,
-        key: &[u8],
-        value: &[u8],
-    ) -> anyhow::Result<()> {
-        let state = StateRecord::Data {
-            account_id: contract_id.to_owned(),
+        account_id: AccountId,
+        public_key: crate::types::PublicKey,
+    ) -> PatchAccessKeyTransaction {
+        PatchAccessKeyTransaction::new(self, account_id, public_key)
+    }
+
+    // shall we expose convenience patch methods here for consistent API?
+
+    pub(crate) async fn fast_forward(&self, delta_height: u64) -> anyhow::Result<()> {
+        // NOTE: RpcSandboxFastForwardResponse is an empty struct with no fields, so don't do anything with it:
+        self.client()
+            // TODO: replace this with the `query` variant when RpcSandboxFastForwardRequest impls Debug
+            .query_nolog(&RpcSandboxFastForwardRequest { delta_height })
+            .await
+            .map_err(|err| anyhow::anyhow!("Failed to fast forward: {:?}", err))?;
+
+        Ok(())
+    }
+}
+
+//todo: review naming
+#[must_use = "don't forget to .apply() this `SandboxPatchStateBuilder`"]
+pub struct PatchStateTransaction<'s> {
+    sandbox: &'s Sandbox,
+    account_id: AccountId,
+    records: Vec<StateRecord>,
+}
+
+impl<'s> PatchStateTransaction<'s> {
+    pub fn new(sandbox: &'s Sandbox, account_id: AccountId) -> Self {
+        PatchStateTransaction {
+            sandbox,
+            account_id,
+            records: Vec::with_capacity(4),
+        }
+    }
+
+    pub fn data(mut self, key: &[u8], value: &[u8]) -> Self {
+        let data = StateRecord::Data {
+            account_id: self.account_id.clone(),
             data_key: key.to_vec(),
             value: value.to_vec(),
         };
-        let records = vec![state];
 
+        self.records.push(data);
+        self
+    }
+
+    pub fn data_multiple(
+        mut self,
+        kvs: impl IntoIterator<Item = (&'s [u8], &'s [u8])>,
+    ) -> Self {
+        self.extend(kvs);
+        self
+    }
+
+
+    pub async fn transact(self) -> anyhow::Result<()> {
+        let records = self.records;
         // NOTE: RpcSandboxPatchStateResponse is an empty struct with no fields, so don't do anything with it:
         let _patch_resp = self
+            .sandbox
             .client()
             .query(&RpcSandboxPatchStateRequest { records })
             .await
             .map_err(|err| anyhow::anyhow!("Failed to patch state: {:?}", err))?;
 
         Ok(())
     }
+}
 
-    pub(crate) async fn fast_forward(&self, delta_height: u64) -> anyhow::Result<()> {
-        // NOTE: RpcSandboxFastForwardResponse is an empty struct with no fields, so don't do anything with it:
-        self.client()
-            // TODO: replace this with the `query` variant when RpcSandboxFastForwardRequest impls Debug
-            .query_nolog(&RpcSandboxFastForwardRequest { delta_height })
+impl<'s> std::iter::Extend<(&'s [u8], &'s [u8])> for PatchStateTransaction<'s>{
+    fn extend<T: IntoIterator<Item = (&'s [u8], &'s [u8])>>(&mut self, iter: T) {
+        let Self {
+            ref mut records,
+            ref account_id,
+            ..
+        } = self;
+        records.extend(iter.into_iter().map(|(key, value)| StateRecord::Data {
+            account_id: account_id.clone(),
+            data_key: key.to_vec(),
+            value: value.to_vec(),
+        }));
+    }
+}
+
+#[must_use = "don't forget to .apply() this `SandboxPatchStateAccountBuilder`"]
+pub struct PatchStateAccountTransaction<'s> {
+    sandbox: &'s Sandbox,
+    account_id: AccountId,
+    amount: Option<Balance>,
+    locked: Option<Balance>,
+    code_hash: Option<CryptoHash>,
+    storage_usage: Option<StorageUsage>,
+}
+
+impl<'s> PatchStateAccountTransaction<'s> {
+    pub const fn new(sandbox: &'s Sandbox, account_id: AccountId) -> Self {
+        Self {
+            sandbox,
+            account_id,
+            amount: None,
+            locked: None,
+            code_hash: None,
+            storage_usage: None,
+        }
+    }
+
+    pub const fn amount(mut self, amount: Balance) -> Self {
+        self.amount = Some(amount);
+        self
+    }
+
+    pub const fn locked(mut self, locked: Balance) -> Self {
+        self.locked = Some(locked);
+        self
+    }
+
+    pub const fn code_hash(mut self, code_hash: CryptoHash) -> Self {
+        self.code_hash = Some(code_hash);
+        self
+    }
+
+    pub const fn storage_usage(mut self, storage_usage: StorageUsage) -> Self {
+        self.storage_usage = Some(storage_usage);
+        self
+    }
+
+    pub async fn apply(self) -> anyhow::Result<()> {
+        let account_view = self
+            .sandbox
+            .client()
+            .view_account(self.account_id.clone(), None);
+
+        let AccountView {
+            amount: previous_amount,
+            locked: previous_locked,
+            code_hash: previous_code_hash,
+            storage_usage: previous_storage_usage,
+            ..
+        } = account_view
             .await
-            .map_err(|err| anyhow::anyhow!("Failed to fast forward: {:?}", err))?;
+            .map_err(|err| anyhow::anyhow!("Failed to read account: {:?}", err))?;
+
+        let account = StateRecord::Account {
+            account_id: self.account_id.clone(),
+            account: near_primitives::account::Account::new(
+                self.amount.unwrap_or(previous_amount),
+                self.locked.unwrap_or(previous_locked),
+                self.code_hash.unwrap_or(previous_code_hash),
+                self.storage_usage.unwrap_or(previous_storage_usage),
+            ),
+        };
+
+        let records = vec![account];
+
+        // NOTE: RpcSandboxPatchStateResponse is an empty struct with no fields, so don't do anything with it:
+        let _patch_resp = self
+            .sandbox
+            .client()
+            .query(&RpcSandboxPatchStateRequest { records })
+            .await
+            .map_err(|err| anyhow::anyhow!("Failed to patch state: {:?}", err))?;
+
+        Ok(())
+    }
+}
+
+#[must_use = "don't forget to .apply() this `SandboxPatchStateAccountBuilder`"]
+pub struct PatchAccessKeyTransaction<'s> {
+    sandbox: &'s Sandbox,
+    account_id: AccountId,
+    public_key: crate::types::PublicKey,
+    nonce: Nonce,
+}
+
+impl<'s> PatchAccessKeyTransaction<'s> {
+    pub const fn new(
+        sandbox: &'s Sandbox,
+        account_id: AccountId,
+        public_key: crate::types::PublicKey,
+    ) -> Self {
+        Self {
+            sandbox,
+            account_id,
+            public_key,
+            nonce: 0,
+        }
+    }
+
+    pub const fn nonce(mut self, nonce: Nonce) -> Self {
+        self.nonce = nonce;
+        self
+    }
+
+    pub async fn full_access(self) -> anyhow::Result<()> {
+        let mut access_key = near_primitives::account::AccessKey::full_access();
+        access_key.nonce = self.nonce;
+        let access_key = StateRecord::AccessKey {
+            account_id: self.account_id,
+            public_key: self.public_key.into(),
+            access_key,
+        };
+
+        let records = vec![access_key];
+
+        // NOTE: RpcSandboxPatchStateResponse is an empty struct with no fields, so don't do anything with it:
+        let _patch_resp = self
+            .sandbox
+            .client()
+            .query(&RpcSandboxPatchStateRequest { records })
+            .await
+            .map_err(|err| anyhow::anyhow!("Failed to patch state: {:?}", err))?;
+
+        Ok(())
+    }
+
+    pub async fn function_call_access(
+        self,
+        receiver_id: &AccountId,
+        method_names: &[&str],
+        allowance: Option<Balance>,
+    ) -> anyhow::Result<()> {
+        let mut access_key: near_primitives::account::AccessKey =
+            crate::types::AccessKey::function_call_access(receiver_id, method_names, allowance)
+                .into();
+        access_key.nonce = self.nonce;
+        let access_key = StateRecord::AccessKey {
+            account_id: self.account_id,
+            public_key: self.public_key.into(),
+            access_key,
+        };
+
+        let records = vec![access_key];
+
+        // NOTE: RpcSandboxPatchStateResponse is an empty struct with no fields, so don't do anything with it:
+        let _patch_resp = self
+            .sandbox
+            .client()
+            .query(&RpcSandboxPatchStateRequest { records })
+            .await
+            .map_err(|err| anyhow::anyhow!("Failed to patch state: {:?}", err))?;
 
         Ok(())
     }

workspaces/src/types/mod.rs

@@ -28,6 +28,8 @@ pub type Balance = u128;
 /// Height of a specific block
 pub type BlockHeight = u64;
 
+/// StorageUsage is used to count the amount of storage used by a contract.
+pub type StorageUsage = u64;
 /// Key types supported for either a [`SecretKey`] or [`PublicKey`]
 #[derive(Debug, Copy, Clone, Serialize, Deserialize)]
 #[non_exhaustive]