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,6 +1,9 @@
 fn main() {
     let doc_build = cfg!(doc) || std::env::var("DOCS_RS").is_ok();
     if !doc_build && cfg!(feature = "install") {
-        near_sandbox_utils::install().expect("Could not install sandbox");
+        // using unwrap because all the useful error messages are hidden inside
+        near_sandbox_utils::ensure_sandbox_bin().unwrap();
+        // previously the next line was causing near sandbox to be installed every time cargo build/check was called
+        // near_sandbox_utils::install().expect("Could not install sandbox");
     }
 }

workspaces/src/network/mod.rs

@@ -10,13 +10,15 @@ mod server;
 mod testnet;
 pub(crate) mod variants;
 
+pub(crate) use sandbox::SandboxPatchStateAccountBuilder; //not needed directly outside of the crate
+pub(crate) use sandbox::SandboxPatchStateBuilder; //not needed directly outside of the crate
 pub(crate) use variants::DEV_ACCOUNT_SEED;
 
-pub use self::betanet::Betanet;
-pub use self::info::Info;
-pub use self::mainnet::Mainnet;
-pub use self::sandbox::Sandbox;
-pub use self::testnet::Testnet;
-pub use self::variants::{
+pub use betanet::Betanet;
+pub use info::Info;
+pub use mainnet::Mainnet;
+pub use sandbox::Sandbox;
+pub use testnet::Testnet;
+pub use variants::{
     AllowDevAccountCreation, DevAccountDeployer, NetworkClient, NetworkInfo, TopLevelAccountCreator,
 };

workspaces/src/network/sandbox.rs

@@ -4,15 +4,20 @@ 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::account::AccessKey;
+use near_primitives::hash::CryptoHash;
 use near_primitives::state_record::StateRecord;
+use near_primitives::types::StorageUsage;
+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, PublicKey, SecretKey};
 use crate::{Account, Contract, Network, Worker};
 
 // Constant taken from nearcore crate to avoid dependency
@@ -135,36 +140,174 @@ impl Sandbox {
         ImportContractTransaction::new(id.to_owned(), worker.client(), self.client())
     }
 
-    pub(crate) async fn patch_state(
-        &self,
-        contract_id: &AccountId,
-        key: &[u8],
-        value: &[u8],
-    ) -> anyhow::Result<()> {
-        let state = StateRecord::Data {
-            account_id: contract_id.to_owned(),
-            data_key: key.to_vec(),
-            value: value.to_vec(),
+    pub(crate) fn patch_state(&self, account_id: AccountId) -> SandboxPatchStateBuilder {
+        SandboxPatchStateBuilder::new(self, account_id)
+    }
+
+    pub(crate) fn patch_account(&self, account_id: AccountId) -> SandboxPatchStateAccountBuilder {
+        SandboxPatchStateAccountBuilder::new(self, account_id)
+    }
+
+    // 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 SandboxPatchStateBuilder<'s> {
+    sandbox: &'s Sandbox,
+    account_id: AccountId,
+    records: Vec<StateRecord>,
+}
+
+impl<'s> SandboxPatchStateBuilder<'s> {
+    pub fn new(sandbox: &'s Sandbox, account_id: AccountId) -> Self {
+        SandboxPatchStateBuilder {
+            sandbox,
+            account_id,
+            records: Vec::with_capacity(4),
+        }
+    }
+
+    pub fn data(mut self, key: impl Into<Vec<u8>>, value: impl Into<Vec<u8>>) -> Self {
+        let data = StateRecord::Data {
+            account_id: self.account_id.clone(),
+            data_key: key.into(),
+            value: value.into(),
         };
-        let records = vec![state];
 
+        self.records.push(data);
+        self
+    }
+
+    pub fn data_multiple(
+        mut self,
+        kvs: impl IntoIterator<Item = (impl Into<Vec<u8>>, impl Into<Vec<u8>>)>,
+    ) -> Self {
+        let Self {
+            ref mut records,
+            ref account_id,
+            ..
+        } = self;
+        records.extend(kvs.into_iter().map(|(key, value)| StateRecord::Data {
+            account_id: account_id.clone(),
+            data_key: key.into(),
+            value: value.into(),
+        }));
+        self
+    }
+
+    pub fn access_key(mut self, public_key: &PublicKey, access_key: &AccessKey) -> Self {
+        let access_key = StateRecord::AccessKey {
+            account_id: self.account_id.clone(),
+            public_key: public_key.0.clone(),
+            access_key: access_key.clone(),
+        };
+        self.records.push(access_key);
+        self
+    }
+
+    pub async fn apply(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(())
     }
+}
+
+#[must_use = "don't forget to .apply() this `SandboxPatchStateAccountBuilder`"]
+pub struct SandboxPatchStateAccountBuilder<'s> {
+    sandbox: &'s Sandbox,
+    account_id: AccountId,
+    amount: Option<Balance>,
+    locked: Option<Balance>,
+    code_hash: Option<CryptoHash>,
+    storage_usage: Option<StorageUsage>,
+}
+
+impl<'s> SandboxPatchStateAccountBuilder<'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(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 })
+    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(())
     }

workspaces/src/worker/impls.rs

@@ -1,13 +1,14 @@
 use crate::network::{AllowDevAccountCreation, NetworkClient, NetworkInfo, TopLevelAccountCreator};
-use crate::network::{Info, Sandbox};
+use crate::network::{Info, Sandbox, SandboxPatchStateAccountBuilder, SandboxPatchStateBuilder};
 use crate::result::{CallExecution, CallExecutionDetails, ViewResultDetails};
 use crate::rpc::client::{Client, DEFAULT_CALL_DEPOSIT, DEFAULT_CALL_FN_GAS};
 use crate::rpc::patch::ImportContractTransaction;
-use crate::types::{AccountId, Gas, InMemorySigner, SecretKey};
+use crate::types::{AccountId, Gas, InMemorySigner, PublicKey, SecretKey};
 use crate::worker::Worker;
 use crate::{Account, Block, Contract};
 use crate::{AccountDetails, Network};
 use async_trait::async_trait;
+use near_primitives::account::AccessKey;
 use near_primitives::types::Balance;
 use std::collections::HashMap;
 
@@ -174,16 +175,63 @@ impl Worker<Sandbox> {
         self.workspace.import_contract(id, worker)
     }
 
+    /// Patch state into the sandbox network, using builder pattern. This will allow us to set
+    /// state that we have acquired in some manner. This allows us to test random cases that
+    /// are hard to come up naturally as state evolves.
+    pub fn patch_state_builder(&self, account_id: &AccountId) -> SandboxPatchStateBuilder {
+        self.workspace.patch_state(account_id.clone())
+    }
+
+    /// Patch account state using builder pattern
+    pub fn patch_account(&self, account_id: &AccountId) -> SandboxPatchStateAccountBuilder {
+        self.workspace.patch_account(account_id.clone())
+    }
+
     /// Patch state into the sandbox network, given a key and value. This will allow us to set
     /// state that we have acquired in some manner. This allows us to test random cases that
     /// are hard to come up naturally as state evolves.
     pub async fn patch_state(
         &self,
         contract_id: &AccountId,
-        key: &[u8],
-        value: &[u8],
+        key: impl Into<Vec<u8>>,
+        value: impl Into<Vec<u8>>,
+    ) -> anyhow::Result<()> {
+        self.workspace
+            .patch_state(contract_id.clone())
+            .data(key, value)
+            .apply()
+            .await?;
+        Ok(())
+    }
+
+    //todo: add more patch state methods like the previous one
+
+    /// Patch state into the sandbox network. Same as `patch_state` but accepts a sequence of key value pairs
+    pub async fn patch_state_multiple(
+        &self,
+        account_id: &AccountId,
+        kvs: impl IntoIterator<Item = (impl Into<Vec<u8>>, impl Into<Vec<u8>>)>,
+    ) -> anyhow::Result<()> {
+        self.workspace
+            .patch_state(account_id.clone())
+            .data_multiple(kvs)
+            .apply()
+            .await?;
+        Ok(())
+    }
+
+    pub async fn patch_access_key(
+        &self,
+        account_id: &AccountId,
+        public_key: &PublicKey,
+        access_key: &AccessKey,
     ) -> anyhow::Result<()> {
-        self.workspace.patch_state(contract_id, key, value).await
+        self.workspace
+            .patch_state(account_id.clone())
+            .access_key(public_key, access_key)
+            .apply()
+            .await?;
+        Ok(())
     }
 
     /// Fast forward to a point in the future. The delta block height is supplied to tell the