Skip to content

Commit

Permalink
contract: refac + clean up (#71)
Browse files Browse the repository at this point in the history 
* clean up

* clean up + refac

* fix + refac tests

* fix test

* fix test

* fix + fmt
  • Loading branch information
@manlikeHB
manlikeHB authored Feb 18, 2025
1 parent 1db28ee commit f2dc7a3
Show file tree
Hide file tree
Showing 6 changed files with 169 additions and 143 deletions.
11 changes: 6 additions & 5 deletions contracts/src/escrow/escrow.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ pub mod Escrow {
struct Storage {
strk_dispatcher: IERC20Dispatcher,
user_balance: Map::<ContractAddress, u256>,
wager_stake: Map::<u64, u256>, // wager_id -> total stake
#[substorage(v0)]
ownable: OwnableComponent::Storage,
#[substorage(v0)]
Expand Down Expand Up @@ -81,7 +82,6 @@ pub mod Escrow {

#[abi(embed_v0)]
impl EscrowImpl of IEscrow<ContractState> {
//TODO: Add access control and restrict to wager contract
fn deposit_to_wallet(ref self: ContractState, from: ContractAddress, amount: u256) {
self.accesscontrol.assert_only_role(WAGER_ROLE);

Expand All @@ -92,12 +92,12 @@ pub mod Escrow {
let strk_dispatcher = self.strk_dispatcher.read();

// transfers funds to escrow
strk_dispatcher.transfer_from(from, get_contract_address(), amount);
assert(strk_dispatcher.balance_of(from) >= amount, 'Insufficient balance');
self.user_balance.entry(from).write(amount + self.get_balance(from));
strk_dispatcher.transfer_from(from, get_contract_address(), amount);
self.emit(DepositEvent { from, amount });
}

//TODO: restrict to wager contract
fn withdraw_from_wallet(ref self: ContractState, to: ContractAddress, amount: u256) {
self.accesscontrol.assert_only_role(WAGER_ROLE);
let strk_dispatcher = self.strk_dispatcher.read();
Expand All @@ -116,11 +116,12 @@ pub mod Escrow {
self.emit(WithdrawEvent { to, amount });
}

//TODO: restrict to wager contract
fn get_balance(self: @ContractState, address: ContractAddress) -> u256 {
self.accesscontrol.assert_only_role(WAGER_ROLE);
self.user_balance.entry(address).read()
}
//TODO: stake amount?

//TODO
fn fund_wager(self: @ContractState, wager_id: u64, amount: u256) {}
}
}
2 changes: 2 additions & 0 deletions contracts/src/escrow/interface.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,4 +5,6 @@ pub trait IEscrow<TContractState> {
fn deposit_to_wallet(ref self: TContractState, from: ContractAddress, amount: u256);
fn get_balance(self: @TContractState, address: ContractAddress) -> u256;
fn withdraw_from_wallet(ref self: TContractState, to: ContractAddress, amount: u256);
//TODO
fn fund_wager(self: @TContractState, wager_id: u64, amount: u256) {}
}
150 changes: 89 additions & 61 deletions contracts/src/tests/test_escrow.cairo
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,16 +4,17 @@ use starknet::{testing, contract_address_const};
use openzeppelin::token::erc20::interface::{IERC20Dispatcher, IERC20DispatcherTrait};

use contracts::escrow::interface::{IEscrowDispatcher, IEscrowDispatcherTrait};
use contracts::tests::utils::{deploy_mock_erc20, OWNER, BOB, WAGER_ADDRESS, deploy_escrow};
use contracts::tests::utils::{deploy_mock_erc20, OWNER, BOB, deploy_escrow, setup};

use snforge_std::{
declare, ContractClassTrait, DeclareResultTrait, start_cheat_caller_address,
stop_cheat_caller_address, spy_events, EventSpyAssertionsTrait,
};

#[test]
fn test_deposit_to_wallet() {
let (escrow, strk_dispatcher) = deploy_escrow(WAGER_ADDRESS());
#[should_panic(expected: ('Caller is missing role',))]
fn test_deposit_to_wallet_unauthorized() {
let (wager, escrow, strk_dispatcher) = setup();

let amount = 50_u256;

Expand All @@ -22,21 +23,34 @@ fn test_deposit_to_wallet() {
strk_dispatcher.approve(escrow.contract_address, amount);
stop_cheat_caller_address(strk_dispatcher.contract_address);

// Deposit to wallet (caller is Wager Contract, which has WAGER_ROLE)
start_cheat_caller_address(escrow.contract_address, WAGER_ADDRESS()); // Simulate Wager Contract
// Attempt to deposit to wallet from an unauthorized address (BOB)
start_cheat_caller_address(escrow.contract_address, BOB()); // Simulate unauthorized caller
escrow.deposit_to_wallet(OWNER(), amount);
stop_cheat_caller_address(escrow.contract_address);
}

assert(strk_dispatcher.balance_of(escrow.contract_address) == amount, 'wrong amount');
start_cheat_caller_address(escrow.contract_address, WAGER_ADDRESS());
assert(escrow.get_balance(OWNER()) == amount, 'wrong balance');
#[test]
#[should_panic(expected: ('Insufficient balance',))]
fn test_deposit_to_wallet_insufficient_balance() {
let (wager, escrow, strk_dispatcher) = setup();

let amount = 50_u256;

// Approve escrow to spend
start_cheat_caller_address(strk_dispatcher.contract_address, OWNER());
strk_dispatcher.approve(escrow.contract_address, amount);
stop_cheat_caller_address(strk_dispatcher.contract_address);

// Deposit to wallet (caller is Wager Contract, which has WAGER_ROLE)
start_cheat_caller_address(
escrow.contract_address, wager.contract_address
); // Simulate Wager Contract
escrow.deposit_to_wallet(BOB(), amount);
stop_cheat_caller_address(escrow.contract_address);
}

#[test]
#[should_panic(expected: ('Caller is missing role',))]
fn test_deposit_to_wallet_unauthorized() {
let (escrow, strk_dispatcher) = deploy_escrow(WAGER_ADDRESS());
fn test_deposit_to_wallet_ok() {
let (wager, escrow, strk_dispatcher) = setup();

let amount = 50_u256;

Expand All @@ -45,14 +59,20 @@ fn test_deposit_to_wallet_unauthorized() {
strk_dispatcher.approve(escrow.contract_address, amount);
stop_cheat_caller_address(strk_dispatcher.contract_address);

// Attempt to deposit to wallet from an unauthorized address (BOB)
start_cheat_caller_address(escrow.contract_address, BOB()); // Simulate unauthorized caller
// Deposit to wallet (caller is Wager Contract, which has WAGER_ROLE)
start_cheat_caller_address(
escrow.contract_address, wager.contract_address
); // Simulate Wager Contract
escrow.deposit_to_wallet(OWNER(), amount);

assert(strk_dispatcher.balance_of(escrow.contract_address) == amount, 'wrong amount');
assert(escrow.get_balance(OWNER()) == amount, 'wrong balance');
stop_cheat_caller_address(escrow.contract_address);
}

#[test]
fn test_withdraw_from_wallet() {
let (escrow, strk_dispatcher) = deploy_escrow(WAGER_ADDRESS());
let (wager, escrow, strk_dispatcher) = setup();

let initial_balance = 1000_u256;
let withdrawal_amount = 500_u256;
Expand All @@ -63,7 +83,9 @@ fn test_withdraw_from_wallet() {
stop_cheat_caller_address(strk_dispatcher.contract_address);

// Deposit to wallet (caller is Wager Contract, which has WAGER_ROLE)
start_cheat_caller_address(escrow.contract_address, WAGER_ADDRESS()); // Simulate Wager Contract
start_cheat_caller_address(
escrow.contract_address, wager.contract_address
); // Simulate Wager Contract
escrow.deposit_to_wallet(OWNER(), initial_balance);

// Withdraw from wallet (caller is Wager Contract, which has WAGER_ROLE)
Expand All @@ -77,7 +99,7 @@ fn test_withdraw_from_wallet() {
#[test]
#[should_panic(expected: ('Caller is missing role',))]
fn test_withdraw_from_wallet_unauthorized() {
let (escrow, strk_dispatcher) = deploy_escrow(WAGER_ADDRESS());
let (wager, escrow, strk_dispatcher) = setup();

let initial_balance = 1000_u256;
let withdrawal_amount = 500_u256;
Expand All @@ -99,7 +121,7 @@ fn test_withdraw_from_wallet_unauthorized() {

#[test]
fn test_get_balance() {
let (escrow, strk_dispatcher) = deploy_escrow(WAGER_ADDRESS());
let (wager, escrow, strk_dispatcher) = setup();

let amount = 50_u256;

Expand All @@ -109,7 +131,9 @@ fn test_get_balance() {
stop_cheat_caller_address(strk_dispatcher.contract_address);

// Deposit to wallet (caller is Wager Contract, which has WAGER_ROLE)
start_cheat_caller_address(escrow.contract_address, WAGER_ADDRESS()); // Simulate Wager Contract
start_cheat_caller_address(
escrow.contract_address, wager.contract_address
); // Simulate Wager Contract
escrow.deposit_to_wallet(OWNER(), amount);

// Get balance (caller is Wager Contract, which has WAGER_ROLE)
Expand All @@ -119,64 +143,68 @@ fn test_get_balance() {
assert(balance == amount, 'wrong balance');
}

// #[test]
// #[should_panic(expected: ('Invalid address',))]
// fn test_deposit_to_wallet_zero_address() {
// let (escrow, _) = deploy_escrow();
#[test]
#[should_panic(expected: ('Invalid address',))]
fn test_deposit_to_wallet_zero_address() {
let (wager, escrow, strk_dispatcher) = setup();

// // deposit to wallet
// escrow.deposit_to_wallet(0.try_into().unwrap(), 50_u256);
// }
// deposit to wallet
start_cheat_caller_address(escrow.contract_address, wager.contract_address);
escrow.deposit_to_wallet(0.try_into().unwrap(), 50_u256);
}

// #[test]
// #[should_panic(expected: ('Invalid address',))]
// fn test_withdraw_from_wallet_zero_address() {
// let (escrow, _) = deploy_escrow();
#[test]
#[should_panic(expected: ('Invalid address',))]
fn test_withdraw_from_wallet_zero_address() {
let (wager, escrow, strk_dispatcher) = setup();

// // withdraw from wallet
// escrow.withdraw_from_wallet(0.try_into().unwrap(), 500_u256);
// }
// withdraw from wallet
start_cheat_caller_address(escrow.contract_address, wager.contract_address);
escrow.withdraw_from_wallet(0.try_into().unwrap(), 500_u256);
}

// #[test]
// #[should_panic(expected: ('Insufficient funds',))]
// fn test_withdraw_from_wallet_not_enough_balance() {
// let (escrow, strk_dispatcher) = deploy_escrow();
#[test]
#[should_panic(expected: ('Insufficient funds',))]
fn test_withdraw_from_wallet_not_enough_balance() {
let (wager, escrow, strk_dispatcher) = setup();

// let initial_balance = 100_u256;
// let withdrawal_amount = 500_u256;
let initial_balance = 100_u256;
let withdrawal_amount = 500_u256;

// // approve escrow to spend
// start_cheat_caller_address(strk_dispatcher.contract_address, OWNER());
// strk_dispatcher.approve(escrow.contract_address, initial_balance);
// stop_cheat_caller_address(strk_dispatcher.contract_address);
// approve escrow to spend
start_cheat_caller_address(strk_dispatcher.contract_address, OWNER());
strk_dispatcher.approve(escrow.contract_address, initial_balance);
stop_cheat_caller_address(strk_dispatcher.contract_address);

// // deposit to wallet
// escrow.deposit_to_wallet(OWNER(), initial_balance);
// deposit to wallet
start_cheat_caller_address(escrow.contract_address, wager.contract_address);
escrow.deposit_to_wallet(OWNER(), initial_balance);

// escrow.withdraw_from_wallet(OWNER(), withdrawal_amount);
// }
escrow.withdraw_from_wallet(OWNER(), withdrawal_amount);
}

// #[test]
// #[should_panic(expected: ('Amount must be positive',))]
// fn test_amount_must_be_positive() {
// let (escrow, strk_dispatcher) = deploy_escrow();
#[test]
#[should_panic(expected: ('Amount must be positive',))]
fn test_amount_must_be_positive() {
let (wager, escrow, strk_dispatcher) = setup();

// let amount = 0_u256;
let amount = 0_u256;

// // approve escrow to spend
// start_cheat_caller_address(strk_dispatcher.contract_address, OWNER());
// strk_dispatcher.approve(escrow.contract_address, amount);
// stop_cheat_caller_address(strk_dispatcher.contract_address);
// approve escrow to spend
start_cheat_caller_address(strk_dispatcher.contract_address, OWNER());
strk_dispatcher.approve(escrow.contract_address, amount);
stop_cheat_caller_address(strk_dispatcher.contract_address);

// // deposit to wallet
// escrow.deposit_to_wallet(OWNER(), amount);
// assert(escrow.get_balance(OWNER()) > amount, 'amount must be positive');
// }
// deposit to wallet
start_cheat_caller_address(escrow.contract_address, wager.contract_address);
escrow.deposit_to_wallet(OWNER(), amount);
assert(escrow.get_balance(OWNER()) > amount, 'amount must be positive');
}

#[test]
#[should_panic(expected: ('Caller is missing role',))]
fn test_get_balance_unauthorized() {
let (escrow, strk_dispatcher) = deploy_escrow(WAGER_ADDRESS());
let (wager, escrow, strk_dispatcher) = setup();

let amount = 50_u256;

Expand Down
Loading

0 comments on commit f2dc7a3

Please sign in to comment.