instead of account contract, use a simulate function

EkuboProtocol · Jun 19, 2024 · bb592fe · bb592fe
1 parent 3e2eb51
commit bb592fe
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 45 deletions.
diff --git a/src/governor.cairo b/src/governor.cairo
@@ -88,9 +88,12 @@ pub trait IGovernor<TContractState> {
 
     // Replaces the code at this address. This must be self-called via a proposal.
     fn upgrade(ref self: TContractState, class_hash: ClassHash);
+
+    // Can be called only via simulateTransaction to determine the effect of the governor making a set of calls
+    fn simulate(ref self: TContractState, calls: Span<Call>) -> Span<Span<felt252>>;
 }
 
-#[starknet::contract(account)]
+#[starknet::contract]
 pub mod Governor {
     use core::hash::{HashStateTrait, HashStateExTrait};
     use core::num::traits::zero::{Zero};
@@ -99,7 +102,7 @@ pub mod Governor {
     use governance::staker::{IStakerDispatcherTrait};
     use starknet::{
         get_block_timestamp, get_caller_address, get_contract_address,
-        syscalls::{replace_class_syscall}, AccountContract, get_tx_info
+        syscalls::{replace_class_syscall}, get_tx_info
     };
     use super::{
         IStakerDispatcher, ContractAddress, IGovernor, Config, ProposalInfo, Call, ExecutionState,
@@ -447,28 +450,19 @@ pub mod Governor {
 
             replace_class_syscall(class_hash).unwrap();
         }
-    }
 
-    // This implementation exists solely for the purpose of allowing simulation of calls from the governor with the flag to skip validation
-    #[abi(embed_v0)]
-    impl GovernorAccountContractForSimulation of AccountContract<ContractState> {
-        fn __validate_declare__(self: @ContractState, class_hash: felt252) -> felt252 {
-            panic!("Not allowed");
-            0
-        }
-        fn __validate__(ref self: ContractState, calls: Array<Call>) -> felt252 {
-            panic!("Not allowed");
-            0
-        }
-        fn __execute__(ref self: ContractState, mut calls: Array<Call>) -> Array<Span<felt252>> {
-            assert(get_caller_address().is_zero(), 'Invalid caller');
-            let tx_version = get_tx_info().unbox().version.into();
-            assert(tx_version == 1 || tx_version == 3, 'Invalid TX version');
+        fn simulate(ref self: ContractState, mut calls: Span<Call>) -> Span<Span<felt252>> {
+            let tx_version = get_tx_info().unbox().version;
+            assert(
+                tx_version == 0x100000000000000000000000000000001
+                    || tx_version == 0x100000000000000000000000000000003,
+                'Simulation only'
+            );
             let mut results: Array<Span<felt252>> = array![];
             while let Option::Some(call) = calls.pop_front() {
                 results.append(call.execute());
             };
-            results
+            results.span()
         }
     }
 }
diff --git a/src/governor_test.cairo b/src/governor_test.cairo
@@ -11,8 +11,7 @@ use starknet::account::{Call};
 use starknet::{
     get_contract_address, syscalls::deploy_syscall, ClassHash, contract_address_const,
     ContractAddress, get_block_timestamp,
-    testing::{set_block_timestamp, set_contract_address, pop_log, set_version},
-    account::{AccountContractDispatcher, AccountContractDispatcherTrait}
+    testing::{set_block_timestamp, set_contract_address, pop_log, set_version}
 };
 
 fn recipient() -> ContractAddress {
@@ -1016,38 +1015,19 @@ fn test_reconfigure_fails_if_not_self_call() {
 }
 
 #[test]
-#[should_panic(expected: ("Not allowed", 'ENTRYPOINT_FAILED'))]
-fn test_governor_validate_fails() {
+#[should_panic(expected: ('Simulation only', 'ENTRYPOINT_FAILED'))]
+fn test_governor_simulate_fails_tx_version() {
     let (_staker, _token, governor, _config) = setup();
-    AccountContractDispatcher { contract_address: governor.contract_address }
-        .__validate__(array![]);
+    governor.simulate(array![].span());
 }
 
 #[test]
-#[should_panic(expected: ("Not allowed", 'ENTRYPOINT_FAILED'))]
-fn test_governor_validate_declare_fails() {
-    let (_staker, _token, governor, _config) = setup();
-    AccountContractDispatcher { contract_address: governor.contract_address }
-        .__validate_declare__(123);
-}
-
-#[test]
-#[should_panic(expected: ('Invalid caller', 'ENTRYPOINT_FAILED'))]
-fn test_governor_execute_fails_from_non_zero() {
-    let (_staker, _token, governor, _config) = setup();
-    set_contract_address(contract_address_const::<1>());
-    AccountContractDispatcher { contract_address: governor.contract_address }.__execute__(array![]);
-}
-
-#[test]
-#[should_panic(expected: ('Invalid TX version', 'ENTRYPOINT_FAILED'))]
 fn test_governor_execute_fails_invalid_tx_version() {
     let (_staker, _token, governor, _config) = setup();
-    set_version(0);
-    AccountContractDispatcher { contract_address: governor.contract_address }.__execute__(array![]);
+    set_version(0x100000000000000000000000000000001);
+    governor.simulate(array![].span());
 }
 
-
 #[test]
 fn test_reconfigure_succeeds_self_call() {
     let (staker, token, governor, config) = setup();