diff --git a/docs/modules/ROOT/pages/access.adoc b/docs/modules/ROOT/pages/access.adoc
index 90e358838..f51875617 100644
--- a/docs/modules/ROOT/pages/access.adoc
+++ b/docs/modules/ROOT/pages/access.adoc
@@ -48,6 +48,7 @@ mod MyContract {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         OwnableEvent: ownable_component::Event
     }
 
@@ -163,7 +164,9 @@ mod MyContract {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         AccessControlEvent: accesscontrol_component::Event,
+        #[flat]
         SRC5Event: src5_component::Event
     }
 
@@ -240,7 +243,9 @@ mod MyContract {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         AccessControlEvent: accesscontrol_component::Event,
+        #[flat]
         SRC5Event: src5_component::Event
     }
 
diff --git a/docs/modules/ROOT/pages/api/upgrades.adoc b/docs/modules/ROOT/pages/api/upgrades.adoc
index cb6555b3a..5a64892eb 100644
--- a/docs/modules/ROOT/pages/api/upgrades.adoc
+++ b/docs/modules/ROOT/pages/api/upgrades.adoc
@@ -44,7 +44,7 @@ NOTE: This function is usually protected by an xref:access.adoc[Access Control]
 use openzeppelin::upgrades::upgradeable::Upgradeable;
 ```
 
-Upgradeable contract module.
+Upgradeable component.
 
 [.contract-index]
 .Internal Functions
diff --git a/docs/modules/ROOT/pages/introspection.adoc b/docs/modules/ROOT/pages/introspection.adoc
index 8eee8a939..0cc177e5b 100644
--- a/docs/modules/ROOT/pages/introspection.adoc
+++ b/docs/modules/ROOT/pages/introspection.adoc
@@ -59,6 +59,7 @@ mod MyContract {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         SRC5Event: src5_component::Event
     }
 
diff --git a/docs/modules/ROOT/pages/upgrades.adoc b/docs/modules/ROOT/pages/upgrades.adoc
index f09dedb48..cf74d4116 100644
--- a/docs/modules/ROOT/pages/upgrades.adoc
+++ b/docs/modules/ROOT/pages/upgrades.adoc
@@ -34,7 +34,7 @@ fn _upgrade(new_class_hash: ClassHash) {
 
 NOTE: If a contract is deployed without this mechanism, its class hash can still be replaced through {library_calls}.
 
-== `Upgradeable` module
+== `Upgradeable` component
 
 OpenZeppelin Contracts for Cairo provides {upgradeable} to add upgradeability support to your contracts.
 
@@ -49,35 +49,55 @@ NOTE: We will be using the following module to implement the {i_upgradeable} int
 ----
 #[starknet::contract]
 mod UpgradeableContract {
-    use openzeppelin::access::ownable::Ownable;
-    use openzeppelin::upgrades::Upgradeable;
+    use openzeppelin::access::ownable::Ownable as ownable_component;
+    use openzeppelin::upgrades::Upgradeable as upgradeable_component;
     use openzeppelin::upgrades::interface::IUpgradeable;
     use starknet::ClassHash;
     use starknet::ContractAddress;
 
+    component!(path: ownable_component, storage: ownable, event: OwnableEvent);
+    component!(path: upgradeable_component, storage: upgradeable, event: UpgradeableEvent);
+
+    /// Ownable
+    #[abi(embed_v0)]
+    impl OwnableImpl = ownable_component::OwnableImpl<ContractState>;
+    impl OwnableInternalImpl = ownable_component::InternalImpl<ContractState>;
+
+    /// Upgradeable
+    impl UpgradeableInternalImpl = upgradeable_component::InternalImpl<ContractState>;
+
     #[storage]
-    struct Storage {}
+    struct Storage {
+        #[substorage(v0)]
+        ownable: ownable_component::Storage,
+        #[substorage(v0)]
+        upgradeable: upgradeable_component::Storage
+    }
+
+    #[event]
+    #[derive(Drop, starknet::Event)]
+    enum Event {
+        #[flat]
+        OwnableEvent: ownable_component::Event,
+        #[flat]
+        UpgradeableEvent: upgradeable_component::Event
+    }
 
     #[constructor]
-    fn constructor(self: @ContractState, owner: ContractAddress) {
-        let mut unsafe_state = Ownable::unsafe_new_contract_state();
-        Ownable::InternalImpl::initializer(ref unsafe_state, owner);
+    fn constructor(ref self: ContractState, owner: ContractAddress) {
+        self.ownable.initializer(owner);
     }
 
     #[external(v0)]
     impl UpgradeableImpl of IUpgradeable<ContractState> {
         fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
             // This function can only be called by the owner
-            let ownable_state = Ownable::unsafe_new_contract_state();
-            Ownable::InternalImpl::assert_only_owner(@ownable_state);
+            self.ownable.assert_only_owner();
 
             // Replace the class hash upgrading the contract
-            let mut upgradeable_state = Upgradeable::unsafe_new_contract_state();
-            Upgradeable::InternalImpl::_upgrade(ref upgradeable_state, new_class_hash);
+            self.upgradeable._upgrade(new_class_hash);
         }
     }
-
-    (...)
 }
 ----
 
diff --git a/src/tests/mocks.cairo b/src/tests/mocks.cairo
index fd4becb51..cf32182e2 100644
--- a/src/tests/mocks.cairo
+++ b/src/tests/mocks.cairo
@@ -17,5 +17,4 @@ mod snake20_mock;
 mod snake721_mock;
 mod snake_account_mock;
 mod src5_mocks;
-mod upgrades_v1;
-mod upgrades_v2;
+mod upgrades_mocks;
diff --git a/src/tests/mocks/accesscontrol_mocks.cairo b/src/tests/mocks/accesscontrol_mocks.cairo
index e187734ac..b4a304696 100644
--- a/src/tests/mocks/accesscontrol_mocks.cairo
+++ b/src/tests/mocks/accesscontrol_mocks.cairo
@@ -30,7 +30,9 @@ mod DualCaseAccessControlMock {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         AccessControlEvent: accesscontrol_component::Event,
+        #[flat]
         SRC5Event: src5_component::Event
     }
 
@@ -69,7 +71,9 @@ mod SnakeAccessControlMock {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         AccessControlEvent: accesscontrol_component::Event,
+        #[flat]
         SRC5Event: src5_component::Event
     }
 
@@ -109,7 +113,9 @@ mod CamelAccessControlMock {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         AccessControlEvent: accesscontrol_component::Event,
+        #[flat]
         SRC5Event: src5_component::Event
     }
 
diff --git a/src/tests/mocks/initializable_mock.cairo b/src/tests/mocks/initializable_mock.cairo
index 86795a031..494df89e5 100644
--- a/src/tests/mocks/initializable_mock.cairo
+++ b/src/tests/mocks/initializable_mock.cairo
@@ -15,6 +15,7 @@ mod InitializableMock {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         InitializableEvent: initializable_component::Event
     }
 }
diff --git a/src/tests/mocks/ownable_mocks.cairo b/src/tests/mocks/ownable_mocks.cairo
index 81433d063..8c79ee7c9 100644
--- a/src/tests/mocks/ownable_mocks.cairo
+++ b/src/tests/mocks/ownable_mocks.cairo
@@ -21,6 +21,7 @@ mod DualCaseOwnableMock {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         OwnableEvent: ownable_component::Event
     }
 
@@ -50,6 +51,7 @@ mod SnakeOwnableMock {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         OwnableEvent: ownable_component::Event
     }
 
@@ -80,6 +82,7 @@ mod CamelOwnableMock {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         OwnableEvent: ownable_component::Event
     }
 
diff --git a/src/tests/mocks/pausable_mock.cairo b/src/tests/mocks/pausable_mock.cairo
index 5ffc5900a..c1352a9d0 100644
--- a/src/tests/mocks/pausable_mock.cairo
+++ b/src/tests/mocks/pausable_mock.cairo
@@ -17,6 +17,7 @@ mod PausableMock {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         PausableEvent: pausable_component::Event
     }
 }
diff --git a/src/tests/mocks/reentrancy_mock.cairo b/src/tests/mocks/reentrancy_mock.cairo
index fa3d5ac6c..5d1536a0b 100644
--- a/src/tests/mocks/reentrancy_mock.cairo
+++ b/src/tests/mocks/reentrancy_mock.cairo
@@ -41,6 +41,7 @@ mod ReentrancyMock {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         ReentrancyGuardEvent: reentrancy_guard_component::Event
     }
 
diff --git a/src/tests/mocks/src5_mocks.cairo b/src/tests/mocks/src5_mocks.cairo
index d325d39e9..12fa66105 100644
--- a/src/tests/mocks/src5_mocks.cairo
+++ b/src/tests/mocks/src5_mocks.cairo
@@ -21,6 +21,7 @@ mod DualCaseSRC5Mock {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         SRC5Event: src5_component::Event
     }
 }
@@ -43,6 +44,7 @@ mod SnakeSRC5Mock {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         SRC5Event: src5_component::Event
     }
 }
@@ -65,6 +67,7 @@ mod CamelSRC5Mock {
     #[event]
     #[derive(Drop, starknet::Event)]
     enum Event {
+        #[flat]
         SRC5Event: src5_component::Event
     }
 }
diff --git a/src/tests/mocks/upgrades_mocks.cairo b/src/tests/mocks/upgrades_mocks.cairo
new file mode 100644
index 000000000..6cdd9df79
--- /dev/null
+++ b/src/tests/mocks/upgrades_mocks.cairo
@@ -0,0 +1,128 @@
+// These contracts are mocks used to test the core functionality of the upgrade functions.
+// The functions are NOT PROTECTED.
+// DO NOT USE IN PRODUCTION.
+
+use starknet::ClassHash;
+
+#[starknet::interface]
+trait IUpgradesV1<TState> {
+    fn upgrade(ref self: TState, new_class_hash: ClassHash);
+    fn set_value(ref self: TState, val: felt252);
+    fn get_value(self: @TState) -> felt252;
+    fn remove_selector(self: @TState);
+}
+
+trait UpgradesV1Trait<TState> {
+    fn set_value(ref self: TState, val: felt252);
+    fn get_value(self: @TState) -> felt252;
+    fn remove_selector(self: @TState);
+}
+
+#[starknet::contract]
+mod UpgradesV1 {
+    use openzeppelin::upgrades::Upgradeable as upgradeable_component;
+    use starknet::ClassHash;
+    use starknet::ContractAddress;
+
+    component!(path: upgradeable_component, storage: upgradeable, event: UpgradeableEvent);
+
+    impl InternalImpl = upgradeable_component::InternalImpl<ContractState>;
+
+    #[storage]
+    struct Storage {
+        #[substorage(v0)]
+        upgradeable: upgradeable_component::Storage,
+        value: felt252
+    }
+
+    #[event]
+    #[derive(Drop, starknet::Event)]
+    enum Event {
+        #[flat]
+        UpgradeableEvent: upgradeable_component::Event
+    }
+
+    #[external(v0)]
+    fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
+        self.upgradeable._upgrade(new_class_hash);
+    }
+
+    #[external(v0)]
+    impl UpgradesV1Impl of super::UpgradesV1Trait<ContractState> {
+        fn set_value(ref self: ContractState, val: felt252) {
+            self.value.write(val);
+        }
+
+        fn get_value(self: @ContractState) -> felt252 {
+            self.value.read()
+        }
+
+        fn remove_selector(self: @ContractState) {}
+    }
+}
+
+#[starknet::interface]
+trait IUpgradesV2<TState> {
+    fn upgrade(ref self: TState, new_class_hash: ClassHash);
+    fn set_value(ref self: TState, val: felt252);
+    fn set_value2(ref self: TState, val: felt252);
+    fn get_value(self: @TState) -> felt252;
+    fn get_value2(self: @TState) -> felt252;
+}
+
+trait UpgradesV2Trait<TState> {
+    fn set_value(ref self: TState, val: felt252);
+    fn set_value2(ref self: TState, val: felt252);
+    fn get_value(self: @TState) -> felt252;
+    fn get_value2(self: @TState) -> felt252;
+}
+
+#[starknet::contract]
+mod UpgradesV2 {
+    use openzeppelin::upgrades::Upgradeable as upgradeable_component;
+    use starknet::ClassHash;
+    use starknet::ContractAddress;
+
+    component!(path: upgradeable_component, storage: upgradeable, event: UpgradeableEvent);
+
+    impl InternalImpl = upgradeable_component::InternalImpl<ContractState>;
+
+    #[storage]
+    struct Storage {
+        #[substorage(v0)]
+        upgradeable: upgradeable_component::Storage,
+        value: felt252,
+        value2: felt252
+    }
+
+    #[event]
+    #[derive(Drop, starknet::Event)]
+    enum Event {
+        #[flat]
+        UpgradeableEvent: upgradeable_component::Event
+    }
+
+    #[external(v0)]
+    fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
+        self.upgradeable._upgrade(new_class_hash);
+    }
+
+    #[external(v0)]
+    impl UpgradesV2Impl of super::UpgradesV2Trait<ContractState> {
+        fn set_value(ref self: ContractState, val: felt252) {
+            self.value.write(val);
+        }
+
+        fn set_value2(ref self: ContractState, val: felt252) {
+            self.value2.write(val);
+        }
+
+        fn get_value(self: @ContractState) -> felt252 {
+            self.value.read()
+        }
+
+        fn get_value2(self: @ContractState) -> felt252 {
+            self.value2.read()
+        }
+    }
+}
diff --git a/src/tests/mocks/upgrades_v1.cairo b/src/tests/mocks/upgrades_v1.cairo
deleted file mode 100644
index 9391d37dc..000000000
--- a/src/tests/mocks/upgrades_v1.cairo
+++ /dev/null
@@ -1,53 +0,0 @@
-// This contract is a mock used to test the core functionality of the upgrade functions.
-// The functions are NOT PROTECTED.
-// DO NOT USE IN PRODUCTION.
-
-use starknet::ClassHash;
-
-#[starknet::interface]
-trait IUpgradesV1<TState> {
-    fn upgrade(ref self: TState, new_class_hash: ClassHash);
-    fn set_value(ref self: TState, val: felt252);
-    fn get_value(self: @TState) -> felt252;
-    fn remove_selector(self: @TState);
-}
-
-trait UpgradesV1Trait<TState> {
-    fn set_value(ref self: TState, val: felt252);
-    fn get_value(self: @TState) -> felt252;
-    fn remove_selector(self: @TState);
-}
-
-#[starknet::contract]
-mod UpgradesV1 {
-    use openzeppelin::upgrades::interface::IUpgradeable;
-    use openzeppelin::upgrades::upgradeable::Upgradeable;
-    use starknet::ClassHash;
-    use starknet::ContractAddress;
-
-    #[storage]
-    struct Storage {
-        value: felt252
-    }
-
-    #[external(v0)]
-    impl UpgradeableImpl of IUpgradeable<ContractState> {
-        fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
-            let mut unsafe_state = Upgradeable::unsafe_new_contract_state();
-            Upgradeable::InternalImpl::_upgrade(ref unsafe_state, new_class_hash);
-        }
-    }
-
-    #[external(v0)]
-    impl UpgradesV1Impl of super::UpgradesV1Trait<ContractState> {
-        fn set_value(ref self: ContractState, val: felt252) {
-            self.value.write(val);
-        }
-
-        fn get_value(self: @ContractState) -> felt252 {
-            self.value.read()
-        }
-
-        fn remove_selector(self: @ContractState) {}
-    }
-}
diff --git a/src/tests/mocks/upgrades_v2.cairo b/src/tests/mocks/upgrades_v2.cairo
deleted file mode 100644
index 6db611764..000000000
--- a/src/tests/mocks/upgrades_v2.cairo
+++ /dev/null
@@ -1,62 +0,0 @@
-// This contract is a mock used to test the core functionality of the upgrade functions.
-// The functions are NOT PROTECTED.
-// DO NOT USE IN PRODUCTION.
-
-use starknet::ClassHash;
-
-#[starknet::interface]
-trait IUpgradesV2<TState> {
-    fn upgrade(ref self: TState, new_class_hash: ClassHash);
-    fn set_value(ref self: TState, val: felt252);
-    fn set_value2(ref self: TState, val: felt252);
-    fn get_value(self: @TState) -> felt252;
-    fn get_value2(self: @TState) -> felt252;
-}
-
-trait UpgradesV2Trait<TState> {
-    fn set_value(ref self: TState, val: felt252);
-    fn set_value2(ref self: TState, val: felt252);
-    fn get_value(self: @TState) -> felt252;
-    fn get_value2(self: @TState) -> felt252;
-}
-
-#[starknet::contract]
-mod UpgradesV2 {
-    use openzeppelin::upgrades::interface::IUpgradeable;
-    use openzeppelin::upgrades::upgradeable::Upgradeable;
-    use starknet::ClassHash;
-    use starknet::ContractAddress;
-
-    #[storage]
-    struct Storage {
-        value: felt252,
-        value2: felt252,
-    }
-
-    #[external(v0)]
-    impl UpgradeableImpl of IUpgradeable<ContractState> {
-        fn upgrade(ref self: ContractState, new_class_hash: ClassHash) {
-            let mut unsafe_state = Upgradeable::unsafe_new_contract_state();
-            Upgradeable::InternalImpl::_upgrade(ref unsafe_state, new_class_hash);
-        }
-    }
-
-    #[external(v0)]
-    impl UpgradesV2Impl of super::UpgradesV2Trait<ContractState> {
-        fn set_value(ref self: ContractState, val: felt252) {
-            self.value.write(val);
-        }
-
-        fn set_value2(ref self: ContractState, val: felt252) {
-            self.value2.write(val);
-        }
-
-        fn get_value(self: @ContractState) -> felt252 {
-            self.value.read()
-        }
-
-        fn get_value2(self: @ContractState) -> felt252 {
-            self.value2.read()
-        }
-    }
-}
diff --git a/src/tests/upgrades/test_upgradeable.cairo b/src/tests/upgrades/test_upgradeable.cairo
index 51eec586a..b030a3d63 100644
--- a/src/tests/upgrades/test_upgradeable.cairo
+++ b/src/tests/upgrades/test_upgradeable.cairo
@@ -1,15 +1,14 @@
-use openzeppelin::tests::mocks::upgrades_v1::IUpgradesV1Dispatcher;
-use openzeppelin::tests::mocks::upgrades_v1::IUpgradesV1DispatcherTrait;
-use openzeppelin::tests::mocks::upgrades_v1::UpgradesV1;
-use openzeppelin::tests::mocks::upgrades_v2::IUpgradesV2Dispatcher;
-use openzeppelin::tests::mocks::upgrades_v2::IUpgradesV2DispatcherTrait;
-use openzeppelin::tests::mocks::upgrades_v2::UpgradesV2;
+use openzeppelin::tests::mocks::upgrades_mocks::{
+    IUpgradesV1Dispatcher, IUpgradesV1DispatcherTrait, UpgradesV1
+};
+use openzeppelin::tests::mocks::upgrades_mocks::{
+    IUpgradesV2Dispatcher, IUpgradesV2DispatcherTrait, UpgradesV2
+};
 use openzeppelin::tests::utils::constants::{CLASS_HASH_ZERO, ZERO};
 use openzeppelin::tests::utils;
-use openzeppelin::upgrades::upgradeable::Upgradeable::Upgraded;
+use openzeppelin::upgrades::Upgradeable::Upgraded;
 use starknet::ClassHash;
 use starknet::ContractAddress;
-use starknet::Felt252TryIntoClassHash;
 
 const VALUE: felt252 = 123;
 
@@ -47,6 +46,8 @@ fn test_upgraded_event() {
 
     let event = utils::pop_log::<Upgraded>(v1.contract_address).unwrap();
     assert(event.class_hash == V2_CLASS_HASH(), 'Invalid class hash');
+
+    utils::assert_no_events_left(v1.contract_address);
 }
 
 #[test]
diff --git a/src/tests/utils.cairo b/src/tests/utils.cairo
index 2c429c5a3..8e58574e3 100644
--- a/src/tests/utils.cairo
+++ b/src/tests/utils.cairo
@@ -12,10 +12,13 @@ fn deploy(contract_class_hash: felt252, calldata: Array<felt252>) -> ContractAdd
 }
 
 /// Pop the earliest unpopped logged event for the contract as the requested type
-/// and checks there's no more data left on the event, preventing unaccounted params.
+/// and checks there's no more keys or data left on the event, preventing unaccounted params.
 /// This function also removes the first key from the event. This is because indexed
-/// params are set as event keys, but the first event key is always set as the 
+/// params are set as event keys, but the first event key is always set as the
 /// event ID.
+///
+/// This method doesn't currently work for components events that are not flattened,
+/// becuase and extra key is added, pushing the event ID key to the second position.
 fn pop_log<T, impl TDrop: Drop<T>, impl TEvent: starknet::Event<T>>(
     address: ContractAddress
 ) -> Option<T> {
@@ -26,12 +29,13 @@ fn pop_log<T, impl TDrop: Drop<T>, impl TEvent: starknet::Event<T>>(
 
     let ret = starknet::Event::deserialize(ref keys, ref data);
     assert(data.is_empty(), 'Event has extra data');
+    assert(keys.is_empty(), 'Event has extra keys');
     ret
 }
 
 /// Asserts that `expected_keys` exactly matches the indexed keys from `event`.
 /// `expected_keys` must include all indexed event keys for `event` in the order
-/// that they're defined. 
+/// that they're defined.
 fn assert_indexed_keys<T, impl TDrop: Drop<T>, impl TEvent: starknet::Event<T>>(
     event: T, expected_keys: Span<felt252>
 ) {
diff --git a/src/upgrades/interface.cairo b/src/upgrades/interface.cairo
index f8bba7083..c5b72113d 100644
--- a/src/upgrades/interface.cairo
+++ b/src/upgrades/interface.cairo
@@ -1,3 +1,6 @@
+// SPDX-License-Identifier: MIT
+// OpenZeppelin Contracts for Cairo v0.7.0 (upgrades/interface.cairo)
+
 use starknet::ClassHash;
 
 #[starknet::interface]
diff --git a/src/upgrades/upgradeable.cairo b/src/upgrades/upgradeable.cairo
index fcf6c7eec..88886dac3 100644
--- a/src/upgrades/upgradeable.cairo
+++ b/src/upgrades/upgradeable.cairo
@@ -1,4 +1,10 @@
-#[starknet::contract]
+// SPDX-License-Identifier: MIT
+// OpenZeppelin Contracts for Cairo v0.7.0 (upgrades/upgradeable.cairo)
+
+/// # Upgradeable Component
+///
+/// The Upgradeable component provides a mechanism to make a contract upgradeable.
+#[starknet::component]
 mod Upgradeable {
     use starknet::ClassHash;
 
@@ -11,15 +17,22 @@ mod Upgradeable {
         Upgraded: Upgraded
     }
 
+    /// Emitted when the contract is upgraded.
     #[derive(Drop, starknet::Event)]
     struct Upgraded {
         class_hash: ClassHash
     }
 
+    mod Errors {
+        const INVALID_CLASS: felt252 = 'Class hash cannot be zero';
+    }
+
     #[generate_trait]
-    impl InternalImpl of InternalState {
-        fn _upgrade(ref self: ContractState, new_class_hash: ClassHash) {
-            assert(!new_class_hash.is_zero(), 'Class hash cannot be zero');
+    impl InternalImpl<
+        TContractState, +HasComponent<TContractState>
+    > of InternalTrait<TContractState> {
+        fn _upgrade(ref self: ComponentState<TContractState>, new_class_hash: ClassHash) {
+            assert(!new_class_hash.is_zero(), Errors::INVALID_CLASS);
             starknet::replace_class_syscall(new_class_hash).unwrap();
             self.emit(Upgraded { class_hash: new_class_hash });
         }
diff --git a/src/utils.cairo b/src/utils.cairo
index 5396a17f6..ee94b1ce0 100644
--- a/src/utils.cairo
+++ b/src/utils.cairo
@@ -48,15 +48,3 @@ impl Felt252TryIntoBool of TryInto<felt252, bool> {
         }
     }
 }
-
-#[inline(always)]
-fn check_gas() {
-    match gas::withdraw_gas() {
-        Option::Some(_) => {},
-        Option::None(_) => {
-            let mut data = ArrayTrait::new();
-            data.append('Out of gas');
-            panic(data);
-        },
-    }
-}
diff --git a/src/utils/unwrap_and_cast.cairo b/src/utils/unwrap_and_cast.cairo
index f625c08a2..658db27ca 100644
--- a/src/utils/unwrap_and_cast.cairo
+++ b/src/utils/unwrap_and_cast.cairo
@@ -3,7 +3,6 @@
 
 use openzeppelin::utils::Felt252TryIntoBool;
 use starknet::ContractAddress;
-use starknet::Felt252TryIntoContractAddress;
 use starknet::SyscallResult;
 use starknet::SyscallResultTrait;