make the test token more normal

src/airdrop_test.cairo

@@ -1,6 +1,6 @@
-use core::num::traits::zero::{Zero};
 use core::array::{ArrayTrait, SpanTrait};
 use core::hash::{LegacyHash};
+use core::num::traits::zero::{Zero};
 use core::option::{OptionTrait};
 
 use core::result::{Result, ResultTrait};
@@ -367,7 +367,7 @@ fn test_claim_three_claims_one_invalid_via_claim_128() {
 
 fn test_claim_is_valid(root: felt252, claim: Claim, proof: Array<felt252>) {
     let pspan = proof.span();
-    let token = deploy_token(get_contract_address(), claim.amount);
+    let token = deploy_token(get_contract_address(), claim.amount.into());
     let airdrop = deploy(token.contract_address, root);
     token.transfer(airdrop.contract_address, claim.amount.into());
 

src/call_trait_test.cairo

@@ -1,8 +1,8 @@
 use core::array::{Array, ArrayTrait};
 use core::hash::{LegacyHash};
 use core::serde::{Serde};
-use governance::test::test_token::{deploy as deploy_token};
 use governance::call_trait::{CallTrait, HashCall};
+use governance::test::test_token::{deploy as deploy_token};
 use starknet::{contract_address_const, get_contract_address, account::{Call}};
 
 #[test]

src/e2e_test.cairo

@@ -3,7 +3,6 @@ use core::option::{OptionTrait};
 
 use core::result::{Result};
 use core::traits::{TryInto};
-use governance::test::test_token::{deploy as deploy_token};
 use governance::factory::{
     IFactoryDispatcher, IFactoryDispatcherTrait, Factory, DeploymentParameters, DeploymentResult,
 };
@@ -13,6 +12,7 @@ use governance::governor::{Governor, Governor::{to_call_id}};
 use governance::governor::{IGovernorDispatcherTrait};
 use governance::interfaces::erc20::{IERC20Dispatcher, IERC20DispatcherTrait};
 use governance::staker::{Staker, IStakerDispatcherTrait};
+use governance::test::test_token::{deploy as deploy_token};
 use governance::timelock::{Timelock, ITimelockDispatcherTrait, TimelockConfig};
 use starknet::account::{Call};
 use starknet::testing::{set_contract_address, set_block_timestamp, pop_log};

src/factory.cairo

@@ -1,6 +1,6 @@
+use governance::airdrop::{Config as AirdropConfig};
 use governance::airdrop::{IAirdropDispatcher};
 use governance::governor::{Config as GovernorConfig};
-use governance::airdrop::{Config as AirdropConfig};
 use governance::governor::{IGovernorDispatcher};
 use governance::staker::{IStakerDispatcher};
 use governance::timelock::{ITimelockDispatcher, TimelockConfig};

src/factory_test.cairo

@@ -21,11 +21,7 @@ use starknet::{
 pub(crate) fn deploy() -> IFactoryDispatcher {
     let mut constructor_args: Array<felt252> = ArrayTrait::new();
     Serde::serialize(
-        @(
-            Staker::TEST_CLASS_HASH,
-            Governor::TEST_CLASS_HASH,
-            Timelock::TEST_CLASS_HASH
-        ),
+        @(Staker::TEST_CLASS_HASH, Governor::TEST_CLASS_HASH, Timelock::TEST_CLASS_HASH),
         ref constructor_args
     );
 

src/staker.cairo

@@ -5,6 +5,11 @@ pub trait IStaker<TContractState> {
     // Returns the token this staker references
     fn get_token(self: @TContractState) -> ContractAddress;
 
+    // Returns the amount staked from the staker to the delegate
+    fn get_staked(
+        self: @TContractState, staker: ContractAddress, delegate: ContractAddress
+    ) -> u128;
+
     // Transfer the approved amount of token from the caller into this contract and delegates it to the given address
     fn stake(ref self: TContractState, delegate: ContractAddress);
 
@@ -202,6 +207,12 @@ pub mod Staker {
             self.token.read().contract_address
         }
 
+        fn get_staked(
+            self: @ContractState, staker: ContractAddress, delegate: ContractAddress
+        ) -> u128 {
+            self.staked.read((staker, delegate))
+        }
+
         fn stake(ref self: ContractState, delegate: ContractAddress) {
             let from = get_caller_address();
             let this_address = get_contract_address();

src/staker_test.cairo

@@ -3,20 +3,20 @@ use core::num::traits::zero::{Zero};
 use core::option::{OptionTrait};
 use core::result::{Result, ResultTrait};
 use core::traits::{TryInto};
-use governance::test::test_token::{TestToken, deploy as deploy_token};
 use governance::interfaces::erc20::{IERC20Dispatcher, IERC20DispatcherTrait};
 
 use governance::staker::{
     IStakerDispatcher, IStakerDispatcherTrait, Staker,
     Staker::{DelegatedSnapshotStorePacking, DelegatedSnapshot},
 };
+use governance::test::test_token::{TestToken, deploy as deploy_token};
 use starknet::testing::{set_contract_address, set_block_timestamp, pop_log};
 use starknet::{
     get_contract_address, syscalls::deploy_syscall, ClassHash, contract_address_const,
     ContractAddress,
 };
 
-pub(crate) fn setup(amount: u128) -> (IStakerDispatcher, IERC20Dispatcher) {
+pub(crate) fn setup(amount: u256) -> (IStakerDispatcher, IERC20Dispatcher) {
     let token = deploy_token(get_contract_address(), amount);
     let (staker_address, _) = deploy_syscall(
         Staker::TEST_CLASS_HASH.try_into().unwrap(),
@@ -33,13 +33,21 @@ mod stake_withdraw {
         setup, Staker, IStakerDispatcherTrait, IERC20Dispatcher, IERC20DispatcherTrait, pop_log,
         get_contract_address, Zero, TestToken, contract_address_const
     };
-    
+
     #[test]
     fn test_takes_approved_token() {
         let (staker, token) = setup(1000);
 
         token.approve(staker.contract_address, 500);
         staker.stake(contract_address_const::<'delegate'>());
+
+        assert_eq!(
+            staker.get_staked(get_contract_address(), contract_address_const::<'delegate'>()), 500
+        );
+        assert_eq!(staker.get_staked(get_contract_address(), Zero::zero()), 0);
+        assert_eq!(
+            staker.get_staked(contract_address_const::<'delegate'>(), get_contract_address()), 0
+        );
         // pop the transfer from 0 to deployer
         pop_log::<TestToken::Transfer>(token.contract_address).unwrap();
         assert_eq!(
@@ -61,6 +69,24 @@ mod stake_withdraw {
             )
         );
     }
+
+    #[test]
+    #[should_panic(expected: ('ALLOWANCE_OVERFLOW', 'ENTRYPOINT_FAILED'))]
+    fn test_fails_allowance_large() {
+        let (staker, token) = setup(1000);
+
+        token.approve(staker.contract_address, u256 { high: 1, low: 0 });
+        staker.stake(contract_address_const::<'delegate'>());
+    }
+
+    #[test]
+    #[should_panic(expected: ('INSUFFICIENT_TF_BALANCE', 'ENTRYPOINT_FAILED', 'ENTRYPOINT_FAILED'))]
+    fn test_fails_insufficient_balance() {
+        let (staker, token) = setup(1000);
+
+        token.approve(staker.contract_address, 1001);
+        staker.stake(contract_address_const::<'delegate'>());
+    }
 }
 
 #[test]

src/test/test_token.cairo

@@ -1,3 +1,6 @@
+use governance::interfaces::erc20::{IERC20Dispatcher};
+
+use starknet::{ContractAddress, syscalls::{deploy_syscall}};
 #[starknet::contract]
 pub(crate) mod TestToken {
     use core::num::traits::zero::{Zero};
@@ -6,8 +9,8 @@ pub(crate) mod TestToken {
 
     #[storage]
     struct Storage {
-        balances: LegacyMap<ContractAddress, u128>,
-        allowances: LegacyMap<(ContractAddress, ContractAddress), u128>,
+        balances: LegacyMap<ContractAddress, u256>,
+        allowances: LegacyMap<(ContractAddress, ContractAddress), u256>,
     }
 
     #[derive(starknet::Event, PartialEq, Debug, Drop)]
@@ -24,9 +27,9 @@ pub(crate) mod TestToken {
     }
 
     #[constructor]
-    fn constructor(ref self: ContractState, recipient: ContractAddress, amount: u128) {
+    fn constructor(ref self: ContractState, recipient: ContractAddress, amount: u256) {
         self.balances.write(recipient, amount);
-        self.emit(Transfer { from: Zero::zero(), to: recipient, value: amount.into() })
+        self.emit(Transfer { from: Zero::zero(), to: recipient, value: amount })
     }
 
     #[abi(embed_v0)]
@@ -40,11 +43,10 @@ pub(crate) mod TestToken {
             self.allowances.read((owner, spender)).into()
         }
         fn transfer(ref self: ContractState, recipient: ContractAddress, amount: u256) -> bool {
-            let amount_small: u128 = amount.try_into().unwrap();
             let balance = self.balances.read(get_caller_address());
-            assert(balance >= amount_small, 'INSUFFICIENT_TRANSFER_BALANCE');
-            self.balances.write(recipient, self.balances.read(recipient) + amount_small);
-            self.balances.write(get_caller_address(), balance - amount_small);
+            assert(balance >= amount, 'INSUFFICIENT_TRANSFER_BALANCE');
+            self.balances.write(recipient, self.balances.read(recipient) + amount);
+            self.balances.write(get_caller_address(), balance - amount);
             self.emit(Transfer { from: get_caller_address(), to: recipient, value: amount });
             true
         }
@@ -54,14 +56,13 @@ pub(crate) mod TestToken {
             recipient: ContractAddress,
             amount: u256
         ) -> bool {
-            let amount_small: u128 = amount.try_into().unwrap();
             let allowance = self.allowances.read((sender, get_caller_address()));
-            assert(allowance >= amount_small, 'INSUFFICIENT_ALLOWANCE');
+            assert(allowance >= amount, 'INSUFFICIENT_ALLOWANCE');
             let balance = self.balances.read(sender);
-            assert(balance >= amount_small, 'INSUFFICIENT_TF_BALANCE');
-            self.balances.write(recipient, self.balances.read(recipient) + amount_small);
-            self.balances.write(sender, balance - amount_small);
-            self.allowances.write((sender, get_caller_address()), allowance - amount_small);
+            assert(balance >= amount, 'INSUFFICIENT_TF_BALANCE');
+            self.balances.write(recipient, self.balances.read(recipient) + amount);
+            self.balances.write(sender, balance - amount);
+            self.allowances.write((sender, get_caller_address()), allowance - amount);
             self.emit(Transfer { from: sender, to: recipient, value: amount });
             true
         }
@@ -72,11 +73,8 @@ pub(crate) mod TestToken {
     }
 }
 
-use starknet::{ContractAddress, syscalls::{deploy_syscall}};
-use governance::interfaces::erc20::{IERC20Dispatcher};
-
 #[cfg(test)]
-pub(crate) fn deploy(owner: ContractAddress, amount: u128) -> IERC20Dispatcher {
+pub(crate) fn deploy(owner: ContractAddress, amount: u256) -> IERC20Dispatcher {
     let mut constructor_args: Array<felt252> = ArrayTrait::new();
     Serde::serialize(@(owner, amount), ref constructor_args);
 

src/timelock_test.cairo

@@ -1,6 +1,6 @@
 use core::array::{Array, ArrayTrait, SpanTrait};
-use governance::test::test_token::{deploy as deploy_token};
 use governance::interfaces::erc20::{IERC20Dispatcher, IERC20DispatcherTrait};
+use governance::test::test_token::{deploy as deploy_token};
 use governance::timelock::{
     ITimelockDispatcher, ITimelockDispatcherTrait, Timelock, TimelockConfig,
     TimelockConfigStorePacking, ExecutionState, ExecutionStateStorePacking