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,16 @@ mod server;
 mod testnet;
 pub(crate) mod variants;
 
+pub(crate) use sandbox::SandboxPatchAcessKeyBuilder; //not needed directly outside of the crate
+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,19 @@ 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::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, Nonce, SecretKey};
 use crate::{Account, Contract, Network, Worker};
 
 // Constant taken from nearcore crate to avoid dependency
@@ -135,36 +139,261 @@ 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) -> SandboxPatchStateBuilder {
+        SandboxPatchStateBuilder::new(self, account_id)
+    }
+
+    pub(crate) fn patch_account(&self, account_id: AccountId) -> SandboxPatchStateAccountBuilder {
+        SandboxPatchStateAccountBuilder::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(),
-            data_key: key.to_vec(),
-            value: value.to_vec(),
+        account_id: AccountId,
+        public_key: crate::types::PublicKey,
+    ) -> SandboxPatchAcessKeyBuilder {
+        SandboxPatchAcessKeyBuilder::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 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(())
     }
+}
 
-    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 })
+#[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 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 SandboxPatchAcessKeyBuilder<'s> {
+    sandbox: &'s Sandbox,
+    account_id: AccountId,
+    public_key: crate::types::PublicKey,
+    nonce: Nonce,
+}
+
+impl<'s> SandboxPatchAcessKeyBuilder<'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(())
     }