From 6962bbe9f9a6e72c81384e70b1f2db66aa84658d Mon Sep 17 00:00:00 2001
From: immrsd <103599616+immrsd@users.noreply.github.com>
Date: Tue, 19 Nov 2024 19:00:38 +0100
Subject: [PATCH] Apply minor improvements to Upgrades package (#1225)

---
 packages/upgrades/src/tests/test_upgradeable.cairo |  8 +++-----
 packages/upgrades/src/upgradeable.cairo            | 14 ++++++++------
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/packages/upgrades/src/tests/test_upgradeable.cairo b/packages/upgrades/src/tests/test_upgradeable.cairo
index f34967437..fa611cf24 100644
--- a/packages/upgrades/src/tests/test_upgradeable.cairo
+++ b/packages/upgrades/src/tests/test_upgradeable.cairo
@@ -1,6 +1,5 @@
-use openzeppelin_test_common::mocks::upgrades::{
-    IUpgradesV1SafeDispatcher, IUpgradesV1SafeDispatcherTrait
-};
+use openzeppelin_test_common::mocks::upgrades::IUpgradesV1SafeDispatcher;
+use openzeppelin_test_common::mocks::upgrades::IUpgradesV1SafeDispatcherTrait;
 use openzeppelin_test_common::mocks::upgrades::{IUpgradesV1Dispatcher, IUpgradesV1DispatcherTrait};
 use openzeppelin_test_common::mocks::upgrades::{IUpgradesV2Dispatcher, IUpgradesV2DispatcherTrait};
 use openzeppelin_test_common::upgrades::UpgradeableSpyHelpers;
@@ -26,7 +25,7 @@ fn setup_test() -> (IUpgradesV1Dispatcher, ContractClass) {
 //
 
 #[test]
-#[should_panic(expected: ('Class hash cannot be zero',))]
+#[should_panic(expected: 'Class hash cannot be zero')]
 fn test_upgrade_with_class_hash_zero() {
     let (v1, _) = setup_test();
     v1.upgrade(CLASS_HASH_ZERO());
@@ -87,7 +86,6 @@ fn test_remove_selector_fails_in_v2() {
     );
 }
 
-
 //
 // upgrade_and_call
 //
diff --git a/packages/upgrades/src/upgradeable.cairo b/packages/upgrades/src/upgradeable.cairo
index 9d544349e..43810c83c 100644
--- a/packages/upgrades/src/upgradeable.cairo
+++ b/packages/upgrades/src/upgradeable.cairo
@@ -7,9 +7,7 @@
 #[starknet::component]
 pub mod UpgradeableComponent {
     use core::num::traits::Zero;
-    use starknet::ClassHash;
-    use starknet::SyscallResultTrait;
-    use starknet::syscalls::{call_contract_syscall, replace_class_syscall};
+    use starknet::{ClassHash, SyscallResultTrait};
 
     #[storage]
     pub struct Storage {}
@@ -30,6 +28,10 @@ pub mod UpgradeableComponent {
         pub const INVALID_CLASS: felt252 = 'Class hash cannot be zero';
     }
 
+    //
+    // Internal
+    //
+
     #[generate_trait]
     pub impl InternalImpl<
         TContractState, +HasComponent<TContractState>
@@ -42,8 +44,8 @@ pub mod UpgradeableComponent {
         ///
         /// Emits an `Upgraded` event.
         fn upgrade(ref self: ComponentState<TContractState>, new_class_hash: ClassHash) {
-            assert(!new_class_hash.is_zero(), Errors::INVALID_CLASS);
-            replace_class_syscall(new_class_hash).unwrap_syscall();
+            assert(new_class_hash.is_non_zero(), Errors::INVALID_CLASS);
+            starknet::syscalls::replace_class_syscall(new_class_hash).unwrap_syscall();
             self.emit(Upgraded { class_hash: new_class_hash });
         }
 
@@ -68,7 +70,7 @@ pub mod UpgradeableComponent {
             // `call_contract_syscall` is used in order to call `selector` from the new class.
             // See:
             // https://docs.starknet.io/documentation/architecture_and_concepts/Contracts/system-calls-cairo1/#replace_class
-            call_contract_syscall(this, selector, calldata).unwrap_syscall()
+            starknet::syscalls::call_contract_syscall(this, selector, calldata).unwrap_syscall()
         }
     }
 }