BCI-1923: Cairo Test Upgradeable (#326)

* remove ignore

* add mock contracts

* finish test

* remove unused code

* fix mock constructor + define impl of trait

* use upgradeable interface for multisig

.tool-versions

@@ -9,6 +9,7 @@ mockery 2.22.1
 golangci-lint 1.55.0
 actionlint 1.6.12
 shellcheck 0.8.0
+scarb 0.7.0
 
 # Kubernetes
 k3d 5.4.4

contracts/src/access_control/access_controller.cairo

@@ -7,7 +7,7 @@ mod AccessController {
 
     use chainlink::libraries::access_control::{AccessControl, IAccessController};
     use chainlink::libraries::ownable::{Ownable, IOwnable};
-    use chainlink::libraries::upgradeable::Upgradeable;
+    use chainlink::libraries::upgradeable::{Upgradeable, IUpgradeable};
 
     #[storage]
     struct Storage {}
@@ -98,9 +98,11 @@ mod AccessController {
     }
 
     #[external(v0)]
-    fn upgrade(ref self: ContractState, new_impl: ClassHash) {
-        let ownable = Ownable::unsafe_new_contract_state();
-        Ownable::assert_only_owner(@ownable);
-        Upgradeable::upgrade(new_impl);
+    impl UpgradeableImpl of IUpgradeable<ContractState> {
+        fn upgrade(ref self: ContractState, new_impl: ClassHash) {
+            let ownable = Ownable::unsafe_new_contract_state();
+            Ownable::assert_only_owner(@ownable);
+            Upgradeable::upgrade(new_impl);
+        }
     }
 }

contracts/src/libraries.cairo

@@ -2,3 +2,4 @@ mod ownable;
 mod access_control;
 mod token;
 mod upgradeable;
+mod mocks;

contracts/src/libraries/mocks.cairo

@@ -0,0 +1,2 @@
+mod mock_upgradeable;
+mod mock_non_upgradeable;

contracts/src/libraries/mocks/mock_non_upgradeable.cairo

@@ -0,0 +1,20 @@
+#[starknet::interface]
+trait IMockNonUpgradeable<TContractState> {
+    fn bar(self: @TContractState) -> bool;
+}
+
+#[starknet::contract]
+mod MockNonUpgradeable {
+    #[storage]
+    struct Storage {}
+
+    #[constructor]
+    fn constructor(ref self: ContractState) {}
+
+    #[external(v0)]
+    impl MockNonUpgradeableImpl of super::IMockNonUpgradeable<ContractState> {
+        fn bar(self: @ContractState) -> bool {
+            true
+        }
+    }
+}

contracts/src/libraries/mocks/mock_upgradeable.cairo

@@ -0,0 +1,31 @@
+use starknet::class_hash::ClassHash;
+
+#[starknet::interface]
+trait IMockUpgradeable<TContractState> {
+    fn foo(self: @TContractState) -> bool;
+    fn upgrade(ref self: TContractState, new_impl: ClassHash);
+}
+
+#[starknet::contract]
+mod MockUpgradeable {
+    use starknet::class_hash::ClassHash;
+
+    use chainlink::libraries::upgradeable::Upgradeable;
+
+    #[storage]
+    struct Storage {}
+
+    #[constructor]
+    fn constructor(ref self: ContractState) {}
+
+    #[external(v0)]
+    impl MockUpgradeableImpl of super::IMockUpgradeable<ContractState> {
+        fn foo(self: @ContractState) -> bool {
+            true
+        }
+
+        fn upgrade(ref self: ContractState, new_impl: ClassHash) {
+            Upgradeable::upgrade(new_impl)
+        }
+    }
+}

contracts/src/libraries/upgradeable.cairo

@@ -2,6 +2,7 @@ use starknet::class_hash::ClassHash;
 
 #[starknet::interface]
 trait IUpgradeable<TContractState> {
+    // note: any contract that uses this module will have a mutable reference to contract state
     fn upgrade(ref self: TContractState, new_impl: ClassHash);
 }
 

contracts/src/multisig.cairo

@@ -50,7 +50,6 @@ trait IMultisig<TContractState> {
     fn is_executed(self: @TContractState, nonce: u128) -> bool;
     fn get_transaction(self: @TContractState, nonce: u128) -> (Transaction, Array::<felt252>);
     fn type_and_version(self: @TContractState) -> felt252;
-    fn upgrade(ref self: TContractState, new_impl: ClassHash);
     fn submit_transaction(
         ref self: TContractState,
         to: ContractAddress,
@@ -93,7 +92,7 @@ mod Multisig {
     use starknet::storage_write_syscall;
     use starknet::class_hash::ClassHash;
 
-    use chainlink::libraries::upgradeable::Upgradeable;
+    use chainlink::libraries::upgradeable::{Upgradeable, IUpgradeable};
 
     #[event]
     #[derive(Drop, starknet::Event)]
@@ -162,6 +161,14 @@ mod Multisig {
         self._set_threshold(threshold);
     }
 
+    #[external(v0)]
+    impl UpgradeableImpl of IUpgradeable<ContractState> {
+        fn upgrade(ref self: ContractState, new_impl: ClassHash) {
+            self._require_multisig();
+            Upgradeable::upgrade(new_impl)
+        }
+    }
+
     #[external(v0)]
     impl MultisigImpl of super::IMultisig<ContractState> {
         /// Views
@@ -214,11 +221,6 @@ mod Multisig {
 
         /// Externals
 
-        fn upgrade(ref self: ContractState, new_impl: ClassHash) {
-            self._require_multisig();
-            Upgradeable::upgrade(new_impl)
-        }
-
         fn submit_transaction(
             ref self: ContractState,
             to: ContractAddress,

contracts/src/ocr2/aggregator.cairo

@@ -197,7 +197,7 @@ mod Aggregator {
     use chainlink::utils::split_felt;
     use chainlink::libraries::ownable::{Ownable, IOwnable};
     use chainlink::libraries::access_control::AccessControl;
-    use chainlink::libraries::upgradeable::Upgradeable;
+    use chainlink::libraries::upgradeable::{Upgradeable, IUpgradeable};
 
     use openzeppelin::token::erc20::interface::{IERC20, IERC20Dispatcher, IERC20DispatcherTrait};
 
@@ -375,10 +375,12 @@ mod Aggregator {
     // --- Upgradeable ---
 
     #[external(v0)]
-    fn upgrade(ref self: ContractState, new_impl: ClassHash) {
-        let ownable = Ownable::unsafe_new_contract_state();
-        Ownable::assert_only_owner(@ownable);
-        Upgradeable::upgrade(new_impl)
+    impl UpgradeableImpl of IUpgradeable<ContractState> {
+        fn upgrade(ref self: ContractState, new_impl: ClassHash) {
+            let ownable = Ownable::unsafe_new_contract_state();
+            Ownable::assert_only_owner(@ownable);
+            Upgradeable::upgrade(new_impl)
+        }
     }
 
     // --- Ownership ---

contracts/src/ocr2/aggregator_proxy.cairo

@@ -52,7 +52,7 @@ mod AggregatorProxy {
     use chainlink::libraries::ownable::{Ownable, IOwnable};
     use chainlink::libraries::access_control::{AccessControl, IAccessController};
     use chainlink::utils::split_felt;
-    use chainlink::libraries::upgradeable::Upgradeable;
+    use chainlink::libraries::upgradeable::{Upgradeable, IUpgradeable};
 
     const SHIFT: felt252 = 0x100000000000000000000000000000000;
     const MAX_ID: felt252 = 0xffffffffffffffffffffffffffffffff;
@@ -172,12 +172,15 @@ mod AggregatorProxy {
     // -- Upgradeable -- 
 
     #[external(v0)]
-    fn upgrade(ref self: ContractState, new_impl: ClassHash) {
-        let ownable = Ownable::unsafe_new_contract_state();
-        Ownable::assert_only_owner(@ownable);
-        Upgradeable::upgrade(new_impl)
+    impl UpgradeableImpl of IUpgradeable<ContractState> {
+        fn upgrade(ref self: ContractState, new_impl: ClassHash) {
+            let ownable = Ownable::unsafe_new_contract_state();
+            Ownable::assert_only_owner(@ownable);
+            Upgradeable::upgrade(new_impl)
+        }
     }
 
+
     // -- Access Control --
 
     #[external(v0)]

contracts/src/tests/test_access_controller.cairo

@@ -13,6 +13,8 @@ use option::OptionTrait;
 use core::result::ResultTrait;
 
 use chainlink::access_control::access_controller::AccessController;
+use chainlink::access_control::access_controller::AccessController::UpgradeableImpl;
+
 use chainlink::libraries::access_control::{
     IAccessController, IAccessControllerDispatcher, IAccessControllerDispatcherTrait
 };
@@ -34,7 +36,7 @@ fn test_upgrade_not_owner() {
     let sender = setup();
     let mut state = STATE();
 
-    AccessController::upgrade(ref state, class_hash_const::<2>());
+    UpgradeableImpl::upgrade(ref state, class_hash_const::<2>());
 }
 
 #[test]

contracts/src/tests/test_aggregator.cairo

@@ -15,7 +15,9 @@ use core::result::ResultTrait;
 
 use chainlink::ocr2::aggregator::pow;
 use chainlink::ocr2::aggregator::Aggregator;
-use chainlink::ocr2::aggregator::Aggregator::{AggregatorImpl, BillingImpl, PayeeManagementImpl};
+use chainlink::ocr2::aggregator::Aggregator::{
+    AggregatorImpl, BillingImpl, PayeeManagementImpl, UpgradeableImpl
+};
 use chainlink::ocr2::aggregator::Aggregator::BillingConfig;
 use chainlink::ocr2::aggregator::Aggregator::PayeeConfig;
 use chainlink::access_control::access_controller::AccessController;
@@ -155,7 +157,8 @@ fn test_access_control() {
 fn test_upgrade_non_owner() {
     let sender = setup();
     let mut state = STATE();
-    Aggregator::upgrade(ref state, class_hash_const::<123>());
+
+    UpgradeableImpl::upgrade(ref state, class_hash_const::<123>());
 }
 
 // --- Billing tests ---

contracts/src/tests/test_aggregator_proxy.cairo

@@ -17,7 +17,7 @@ use chainlink::ocr2::mocks::mock_aggregator::{
 // use chainlink::ocr2::aggregator::{IAggregator, IAggregatorDispatcher, IAggregatorDispatcherTrait};
 use chainlink::ocr2::aggregator_proxy::AggregatorProxy;
 use chainlink::ocr2::aggregator_proxy::AggregatorProxy::{
-    AggregatorProxyImpl, AggregatorProxyInternal, AccessControllerImpl
+    AggregatorProxyImpl, AggregatorProxyInternal, AccessControllerImpl, UpgradeableImpl
 };
 use chainlink::ocr2::aggregator::Round;
 use chainlink::utils::split_felt;
@@ -95,7 +95,7 @@ fn test_access_control() {
 fn test_upgrade_non_owner() {
     let (_, _, _, _, _) = setup();
     let mut state = STATE();
-    AggregatorProxy::upgrade(ref state, class_hash_const::<123>());
+    UpgradeableImpl::upgrade(ref state, class_hash_const::<123>());
 }
 
 fn test_query_latest_round_data() {

contracts/src/tests/test_link_token.cairo

@@ -13,7 +13,7 @@ use option::OptionTrait;
 use core::result::ResultTrait;
 
 use chainlink::token::link_token::LinkToken;
-use chainlink::token::link_token::LinkToken::{MintableToken, ERC20Impl};
+use chainlink::token::link_token::LinkToken::{MintableToken, ERC20Impl, UpgradeableImpl};
 use chainlink::tests::test_ownable::should_implement_ownable;
 
 // only tests link token specific functionality 
@@ -152,6 +152,6 @@ fn test_upgrade_non_owner() {
     let mut state = STATE();
     LinkToken::constructor(ref state, sender, contract_address_const::<111>());
 
-    LinkToken::upgrade(ref state, class_hash_const::<123>());
+    UpgradeableImpl::upgrade(ref state, class_hash_const::<123>());
 }
 

contracts/src/tests/test_multisig.cairo

@@ -12,7 +12,7 @@ use traits::TryInto;
 
 use chainlink::multisig::assert_unique_values;
 use chainlink::multisig::Multisig;
-use chainlink::multisig::Multisig::MultisigImpl;
+use chainlink::multisig::Multisig::{MultisigImpl, UpgradeableImpl};
 
 #[starknet::contract]
 mod MultisigTest {
@@ -288,7 +288,7 @@ fn test_upgrade_not_multisig() {
     let account = contract_address_const::<777>();
     set_caller_address(account);
 
-    MultisigImpl::upgrade(ref state, class_hash_const::<1>())
+    UpgradeableImpl::upgrade(ref state, class_hash_const::<1>())
 }
 
 #[test]

contracts/src/tests/test_upgradeable.cairo

@@ -4,27 +4,55 @@ use starknet::testing::set_caller_address;
 use starknet::ContractAddress;
 use starknet::contract_address_const;
 use starknet::class_hash::class_hash_const;
+use starknet::syscalls::deploy_syscall;
 
 use chainlink::libraries::upgradeable::Upgradeable;
 use chainlink::libraries::ownable::Ownable;
-
+use chainlink::libraries::mocks::mock_upgradeable::{
+    MockUpgradeable, IMockUpgradeableDispatcher, IMockUpgradeableDispatcherTrait,
+    IMockUpgradeableDispatcherImpl
+};
+use chainlink::libraries::mocks::mock_non_upgradeable::{
+    MockNonUpgradeable, IMockNonUpgradeableDispatcher, IMockNonUpgradeableDispatcherTrait,
+    IMockNonUpgradeableDispatcherImpl
+};
 
 fn setup() -> ContractAddress {
     let account: ContractAddress = contract_address_const::<777>();
     set_caller_address(account);
     account
 }
 
-// replace_class_syscall() not yet supported in tests
 #[test]
-#[ignore]
 #[available_gas(2000000)]
 fn test_upgrade() {
     let sender = setup();
 
+    // doesn't error
     Upgradeable::upgrade(class_hash_const::<1>());
 }
 
+#[test]
+#[available_gas(2000000)]
+fn test_upgrade_and_call() {
+    let sender = setup();
+
+    let calldata = array![];
+    let (contractAddr, _) = deploy_syscall(
+        MockUpgradeable::TEST_CLASS_HASH.try_into().unwrap(), 0, calldata.span(), false
+    )
+        .unwrap();
+    let mockUpgradeable = IMockUpgradeableDispatcher { contract_address: contractAddr };
+    assert(mockUpgradeable.foo() == true, 'should call foo');
+
+    mockUpgradeable.upgrade(MockNonUpgradeable::TEST_CLASS_HASH.try_into().unwrap());
+
+    // now, contract should be different
+    let mockNonUpgradeable = IMockNonUpgradeableDispatcher { contract_address: contractAddr };
+    assert(mockNonUpgradeable.bar() == true, 'should call bar');
+}
+
+
 #[test]
 #[available_gas(2000000)]
 #[should_panic(expected: ('Class hash cannot be zero',))]

contracts/src/token/link_token.cairo

@@ -21,7 +21,7 @@ mod LinkToken {
     use openzeppelin::token::erc20::interface::{IERC20, IERC20Dispatcher, IERC20DispatcherTrait};
     use chainlink::libraries::token::erc677::ERC677;
     use chainlink::libraries::ownable::{Ownable, IOwnable};
-    use chainlink::libraries::upgradeable::Upgradeable;
+    use chainlink::libraries::upgradeable::{Upgradeable, IUpgradeable};
 
     const NAME: felt252 = 'ChainLink Token';
     const SYMBOL: felt252 = 'LINK';
@@ -86,10 +86,12 @@ mod LinkToken {
     //  Upgradeable
     //
     #[external(v0)]
-    fn upgrade(ref self: ContractState, new_impl: ClassHash) {
-        let ownable = Ownable::unsafe_new_contract_state();
-        Ownable::assert_only_owner(@ownable);
-        Upgradeable::upgrade(new_impl)
+    impl UpgradeableImpl of IUpgradeable<ContractState> {
+        fn upgrade(ref self: ContractState, new_impl: ClassHash) {
+            let ownable = Ownable::unsafe_new_contract_state();
+            Ownable::assert_only_owner(@ownable);
+            Upgradeable::upgrade(new_impl)
+        }
     }
 
     //