From f924927810d905ca5c7f7a6ec301ad85cdb27f4c Mon Sep 17 00:00:00 2001
From: Oleksandr Anyshchenko <oleksandr.anyshchenko@aurora.dev>
Date: Tue, 30 Jan 2024 17:00:37 +0000
Subject: [PATCH] chore: use enum for collection prefixes (#4)

---
 src/keys.rs                     | 21 +++++++++++++++++----
 src/lib.rs                      | 12 ++++++------
 src/tests/sdk/mod.rs            | 32 +-------------------------------
 src/tests/workspace/delegate.rs |  4 ++--
 src/types.rs                    |  7 +++----
 5 files changed, 29 insertions(+), 47 deletions(-)

diff --git a/src/keys.rs b/src/keys.rs
index a42f12e..3894731 100644
--- a/src/keys.rs
+++ b/src/keys.rs
@@ -1,4 +1,17 @@
-pub const BLOBS_KEY: &[u8] = b"BLOBS";
-pub const DEPLOYMENTS_KEY: &[u8] = b"DEPLOYMENTS";
-pub const RELEASES_KEY: &[u8] = b"RELEASES";
-pub const LATEST_KEY: &[u8] = b"LATEST";
+use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
+use near_sdk::IntoStorageKey;
+
+/// Prefixes for the `near-sdk` collections that are used in the smart contract.
+#[derive(BorshSerialize, BorshDeserialize)]
+pub enum Prefix {
+    Blobs,
+    Deployments,
+    Releases,
+    LatestRelease,
+}
+
+impl IntoStorageKey for Prefix {
+    fn into_storage_key(self) -> Vec<u8> {
+        self.try_to_vec().unwrap()
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index bc6f69c..a174efd 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -87,10 +87,10 @@ impl AuroraControllerFactory {
     pub fn new(owner_id: AccountId) -> Self {
         assert!(!env::state_exists(), "Already initialized");
         let mut contract = Self {
-            releases: UnorderedMap::new(keys::RELEASES_KEY),
-            blobs: UnorderedMap::new(keys::BLOBS_KEY),
-            deployments: UnorderedMap::new(keys::DEPLOYMENTS_KEY),
-            latest: LazyOption::new(keys::LATEST_KEY, None),
+            releases: UnorderedMap::new(keys::Prefix::Releases),
+            blobs: UnorderedMap::new(keys::Prefix::Blobs),
+            deployments: UnorderedMap::new(keys::Prefix::Deployments),
+            latest: LazyOption::new(keys::Prefix::LatestRelease, None),
         };
 
         contract.owner_set(Some(owner_id));
@@ -134,8 +134,8 @@ impl AuroraControllerFactory {
                 promise.function_call(
                     action.function_name,
                     action.arguments.into(),
-                    action.amount,
-                    action.gas,
+                    action.amount.into(),
+                    Gas(action.gas.0),
                 )
             })
     }
diff --git a/src/tests/sdk/mod.rs b/src/tests/sdk/mod.rs
index c0cd48f..0d74c54 100644
--- a/src/tests/sdk/mod.rs
+++ b/src/tests/sdk/mod.rs
@@ -1,8 +1,7 @@
 use crate::types::ReleaseInfo;
 use crate::AuroraControllerFactory;
 use near_plugins::AccessControllable;
-use near_sdk::serde_json::Value;
-use near_sdk::{AccountId, ONE_NEAR};
+use near_sdk::AccountId;
 
 #[macro_use]
 mod macros;
@@ -96,35 +95,6 @@ fn test_adding_blob_without_adding_hash() {
     contract.add_release_blob();
 }
 
-#[ignore]
-#[test]
-fn test_deploy_new_aurora_contract() {
-    set_env!(
-        predecessor_account_id: owner(),
-        attached_deposit: ONE_NEAR * 15,
-        input: include_bytes!("../../../res/aurora-mainnet-silo-3.4.0.wasm").to_vec(),
-    );
-    let mut contract = AuroraControllerFactory::new(owner());
-    assert!(contract.acl_grant_role("DAO".to_owned(), owner()).unwrap());
-
-    contract.add_release_info(
-        "05ed0185c7348ca8949727d7f0f20870cf24a768407b2870032e6ef3963990a7".to_string(),
-        "3.4.0".parse().unwrap(),
-        true,
-        None,
-        None,
-    );
-    contract.add_release_blob();
-    contract.deploy(
-        new_engine(),
-        "new".to_string(),
-        Value::default(),
-        Some("05ed0185c7348ca8949727d7f0f20870cf24a768407b2870032e6ef3963990a7".to_string()),
-    );
-
-    assert_eq!(contract.get_deployments().len(), 1);
-}
-
 #[test]
 fn test_check_latest_release() {
     set_env!(
diff --git a/src/tests/workspace/delegate.rs b/src/tests/workspace/delegate.rs
index 2d1e03a..4bdb3c9 100644
--- a/src/tests/workspace/delegate.rs
+++ b/src/tests/workspace/delegate.rs
@@ -19,8 +19,8 @@ async fn test_delegate_execution() {
             "actions": vec![FunctionCallArgs {
                 function_name: "set_owner".to_string(),
                 arguments: contract_id.try_to_vec().map(Into::into).unwrap(),
-                amount: 0,
-                gas: Gas::ONE_TERA * 5
+                amount: 0.into(),
+                gas: (Gas::ONE_TERA * 5).0.into()
             }]
         }))
         .max_gas()
diff --git a/src/types.rs b/src/types.rs
index 45126cc..053a828 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -1,7 +1,6 @@
 use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
-use near_sdk::json_types::Base64VecU8;
+use near_sdk::json_types::{Base64VecU8, U128, U64};
 use near_sdk::serde::{Deserialize, Serialize};
-use near_sdk::{Balance, Gas};
 use std::collections::BTreeMap;
 use std::fmt::{Display, Formatter};
 use std::io::Write;
@@ -94,8 +93,8 @@ impl Display for Version {
 pub struct FunctionCallArgs {
     pub function_name: String,
     pub arguments: Base64VecU8,
-    pub amount: Balance,
-    pub gas: Gas,
+    pub amount: U128,
+    pub gas: U64,
 }
 
 #[test]