Apply minor improvements to Upgrades package (#1225)

OpenZeppelin · Nov 19, 2024 · 6962bbe · 6962bbe
1 parent c77b4b1
commit 6962bbe
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 11 deletions.
diff --git 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
@@ -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()
         }
     }
 }