Skip to content

Commit

Permalink
Merge pull request #12 from hoangkianh/main
Browse files Browse the repository at this point in the history 
Implement a function to add ERC-20 tokens
  • Loading branch information
@Dprof-in-tech
Dprof-in-tech authored Dec 19, 2024
2 parents a2020b4 + 3c590c5 commit e66150d
Show file tree
Hide file tree
Showing 3 changed files with 84 additions and 18 deletions.
30 changes: 26 additions & 4 deletions src/Wallet.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ contract Wallet {

mapping(address => Transaction[]) public transactions;
mapping(address => mapping(address => uint256)) token_balance;
mapping(address => bool) public supportedTokens;

struct Transaction {
uint256 amount;
Expand All @@ -28,22 +29,43 @@ contract Wallet {
_;
}

modifier onlySupportedToken(address _token) {
require(supportedTokens[_token], "Token not supported");
_;
}

constructor(address _owner, address _worldID, address _usdt) {
require(msg.sender != address(0), "zero address found");
owner = _owner;
worldID = IWorldID(_worldID);
usdt = IERC20(_usdt);
supportedTokens[_usdt] = true; // Add USDT as a default supported token
}

function createWorldId() external onlyOwner {}

function transfer(address _recipient, uint256 _amount) external onlyVerified(msg.sender) {
function transfer(address _recipient, address _token, uint256 _amount)
external
onlyVerified(msg.sender)
onlySupportedToken(_token)
{
require(_recipient != address(0), "Zero address detected");
require(usdt.balanceOf(msg.sender) >= _amount, "Insufficient balance");
IERC20 token = IERC20(_token);
require(token.balanceOf(msg.sender) >= _amount, "Insufficient balance");
require(_amount > 0, "Transfer amount must be greater than zero");
require(usdt.transferFrom(msg.sender, _recipient, _amount), "Transfer failed");
require(token.transferFrom(msg.sender, _recipient, _amount), "Transfer failed");

recordTransactionHistory(msg.sender, _amount, _token);
}

function addSupportedToken(address _token) external onlyOwner {
require(_token != address(0), "Invalid token address");
supportedTokens[_token] = true;
}

recordTransactionHistory(msg.sender, _amount, address(usdt));
function removeSupportedToken(address _token) external onlyOwner {
require(supportedTokens[_token], "Token not supported");
supportedTokens[_token] = false;
}

//////////////////////////////////////////////
Expand Down
14 changes: 7 additions & 7 deletions test/Transfer.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,7 @@ contract USDTTransferTest is Test {

// Perform the transfer
vm.prank(user1);
transferContract.transfer(user2, transferAmount);
transferContract.transfer(user2, address(mockUSDT), transferAmount);

// Assert balances after transfer
assertEq(mockUSDT.balanceOf(user1), user1BalanceBefore - transferAmount);
Expand All @@ -134,7 +134,7 @@ contract USDTTransferTest is Test {

vm.prank(user1);
vm.expectRevert("Insufficient balance");
transferContract.transfer(user2, transferAmount);
transferContract.transfer(user2, address(mockUSDT), transferAmount);
}

// Test transfer function - Unverified user
Expand All @@ -146,7 +146,7 @@ contract USDTTransferTest is Test {

vm.prank(user1);
vm.expectRevert("User not verified");
transferContract.transfer(user2, transferAmount);
transferContract.transfer(user2, address(mockUSDT), transferAmount);
}

// Test transfer function - Transfer to unverified recipient
Expand All @@ -157,7 +157,7 @@ contract USDTTransferTest is Test {

vm.expectRevert("User not verified");

transferContract.transfer(user2, transferAmount);
transferContract.transfer(user2, address(mockUSDT), transferAmount);
}

// Test transfer function - Multiple consecutive transfers
Expand All @@ -171,7 +171,7 @@ contract USDTTransferTest is Test {

// First transfer
vm.prank(user1);
transferContract.transfer(user2, transferAmount1);
transferContract.transfer(user2, address(mockUSDT), transferAmount1);

// Assert balances after first transfer
assertEq(mockUSDT.balanceOf(user1), user1BalanceBefore - transferAmount1);
Expand All @@ -182,7 +182,7 @@ contract USDTTransferTest is Test {
user2BalanceBefore = mockUSDT.balanceOf(user2);

vm.prank(user1);
transferContract.transfer(user2, transferAmount2);
transferContract.transfer(user2, address(mockUSDT), transferAmount2);

assertEq(mockUSDT.balanceOf(user1), user1BalanceBefore - transferAmount2);
assertEq(mockUSDT.balanceOf(user2), user2BalanceBefore + transferAmount2);
Expand All @@ -195,6 +195,6 @@ contract USDTTransferTest is Test {
vm.prank(user1);

vm.expectRevert("Transfer amount must be greater than zero");
transferContract.transfer(user2, transferAmount);
transferContract.transfer(user2, address(mockUSDT), transferAmount);
}
}
58 changes: 51 additions & 7 deletions test/Wallet.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,15 +16,20 @@ contract WalletTest is Test {
address public owner;
address public user1;
address public user2;
address public nonOwner;
MockERC20 public usdt;
MockERC20 public anotherToken;
MockWorldID public worldID;

/// @notice Set up the test environment before each test
function setUp() public {
owner = address(this);
user1 = address(0x1);
user2 = address(0x2);
nonOwner = address(0x999);
usdt = new MockERC20("USDT", "USDT");
anotherToken = new MockERC20("Another Token", "ATKN");

worldID = new MockWorldID();

factory = new WalletFactory();
Expand All @@ -48,7 +53,7 @@ contract WalletTest is Test {
/// @notice Test recording a single transaction
function testRecordSingleTransaction() public {
vm.prank(user1);
wallet.transfer(user2, 100);
wallet.transfer(user2, address(usdt), 100);

vm.prank(user1);
Wallet.Transaction[] memory history = wallet.getTransactionHistory(user1);
Expand All @@ -60,8 +65,8 @@ contract WalletTest is Test {
/// @notice Test recording multiple transactions
function testRecordMultipleTransactions() public {
vm.startPrank(user1);
wallet.transfer(user2, 100);
wallet.transfer(user2, 200);
wallet.transfer(user2, address(usdt), 100);
wallet.transfer(user2, address(usdt), 200);

Wallet.Transaction[] memory history = wallet.getTransactionHistory(user1);
vm.stopPrank();
Expand All @@ -76,13 +81,13 @@ contract WalletTest is Test {
/// @notice Test recording transactions for different users
function testRecordTransactionsForDifferentUsers() public {
vm.prank(user1);
wallet.transfer(user2, 100);
wallet.transfer(user2, address(usdt), 100);

vm.prank(user1);
wallet.transfer(user2, 50);
wallet.transfer(user2, address(usdt), 50);

vm.prank(user2);
wallet.transfer(user1, 50);
wallet.transfer(user1, address(usdt), 50);

vm.prank(user1);
Wallet.Transaction[] memory user1History = wallet.getTransactionHistory(user1);
Expand All @@ -103,7 +108,7 @@ contract WalletTest is Test {

vm.startPrank(user1);
usdt.approve(address(wallet), largeAmount);
wallet.transfer(user2, largeAmount);
wallet.transfer(user2, address(usdt), largeAmount);

Wallet.Transaction[] memory history = wallet.getTransactionHistory(user1);
vm.stopPrank();
Expand All @@ -112,4 +117,43 @@ contract WalletTest is Test {
assertEq(history[0].amount, largeAmount, "Recorded amount does not match");
assertEq(history[0].token, address(usdt), "Recorded token address does not match");
}

function testAddSupportedTokenByOwner() public {
vm.prank(owner); // Set the caller as the owner
wallet.addSupportedToken(address(anotherToken));

assertTrue(wallet.supportedTokens(address(anotherToken)));
}

function testAddSupportedTokenByNonOwnerReverts() public {
vm.prank(nonOwner);
vm.expectRevert("not owner");
wallet.addSupportedToken(address(anotherToken));
}

function testRemoveSupportedTokenByOwner() public {
vm.startPrank(owner); // Start a transaction as the owner
wallet.addSupportedToken(address(anotherToken));
assertTrue(wallet.supportedTokens(address(anotherToken)));

wallet.removeSupportedToken(address(anotherToken));
assertFalse(wallet.supportedTokens(address(anotherToken)));
vm.stopPrank();
}

function testRemoveSupportedTokenByNonOwnerReverts() public {
vm.startPrank(owner);
wallet.addSupportedToken(address(anotherToken));
vm.stopPrank();

vm.prank(nonOwner);
vm.expectRevert("not owner");
wallet.removeSupportedToken(address(anotherToken));
}

function testRemoveNonSupportedTokenReverts() public {
vm.prank(owner);
vm.expectRevert("Token not supported");
wallet.removeSupportedToken(address(anotherToken));
}
}

0 comments on commit e66150d

Please sign in to comment.