finish test

smartcontractkit · Sep 19, 2023 · af81ebb · af81ebb
1 parent eb350e4
commit af81ebb
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 0 deletions.
diff --git a/contracts/src/libraries/mocks/mock_non_upgradeable.cairo b/contracts/src/libraries/mocks/mock_non_upgradeable.cairo
@@ -11,6 +11,7 @@ mod MockNonUpgradeable {
     #[constructor]
     fn constructor(self: @ContractState) {}
 
+    #[external(v0)]
     impl MockNonUpgradeableImpl of super::IMockNonUpgradeable<ContractState> {
         fn bar(self: @ContractState) -> bool {
             true

diff --git a/contracts/src/libraries/mocks/mock_upgradeable.cairo b/contracts/src/libraries/mocks/mock_upgradeable.cairo
@@ -18,6 +18,7 @@ mod MockUpgradeable {
     #[constructor]
     fn constructor(self: @ContractState) {}
 
+    #[external(v0)]
     impl MockUpgradeableImpl of super::IMockUpgradeable<ContractState> {
         fn foo(self: @ContractState) -> bool {
             true

diff --git a/contracts/src/libraries/upgradeable.cairo b/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);
 }
 

diff --git a/contracts/src/tests/test_upgradeable.cairo b/contracts/src/tests/test_upgradeable.cairo
@@ -4,9 +4,22 @@ 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 STATE() -> MockUpgradeable::ContractState {
+    MockUpgradeable::contract_state_for_testing()
+}
 
 // Some tests are still written in TS due to missing features in cairo-test
 
@@ -21,9 +34,31 @@ fn setup() -> ContractAddress {
 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',))]