fix: comments from PR

OriumNetwork · Sep 19, 2024 · 432328e · 432328e
1 parent 4c239e1
commit 432328e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 8 deletions.
diff --git a/contracts/ERC20Splitter.sol b/contracts/ERC20Splitter.sol
@@ -60,8 +60,7 @@ contract ERC20Splitter is ReentrancyGuard {
     /// @notice Withdraw all tokens that the caller is entitled to.
     /// Tokens are automatically determined based on previous deposits.
     function withdraw() external nonReentrant {
-        address recipient = msg.sender;
-        address[] storage senderTokens = userTokens[recipient];
+        address[] storage senderTokens = userTokens[msg.sender];
 
         if (senderTokens.length == 0) {
             return;
@@ -71,25 +70,25 @@ contract ERC20Splitter is ReentrancyGuard {
 
         for (uint256 i = 0; i < senderTokens.length; i++) {
             address tokenAddress = senderTokens[i];
-            uint256 amount = balances[tokenAddress][recipient];
+            uint256 amount = balances[tokenAddress][msg.sender];
 
-            balances[tokenAddress][recipient] = 0;
+            delete balances[tokenAddress][msg.sender];
 
             if (tokenAddress == address(0)) {
                 payable(msg.sender).transfer(amount);
             } else {
                 require(
-                    IERC20(tokenAddress).transferFrom(address(this), recipient, amount),
+                    IERC20(tokenAddress).transferFrom(address(this), msg.sender, amount),
                     'ERC20Splitter: TransferFrom failed'
                 );
             }
 
             withdrawnAmounts[i] = amount;
         }
 
-        emit Withdraw(recipient, userTokens[recipient], withdrawnAmounts);
+        emit Withdraw(msg.sender, userTokens[msg.sender], withdrawnAmounts);
 
-        delete userTokens[recipient];
+        delete userTokens[msg.sender];
     }
 
     /** Internal Functions **/

diff --git a/test/SplitterContract.test.ts b/test/SplitterContract.test.ts
@@ -148,7 +148,7 @@ describe('ERC20Splitter', () => {
         expect(await splitter.balances(mockERC20_4.getAddress(), recipient4.address)).to.equal(ethers.parseEther('100'))
       })
 
-      it.only('Should deposit four ERC20 tokens and split them between recipients', async () => {
+      it('Should deposit four ERC20 tokens and split them between recipients', async () => {
         const tokenAmounts = [
           ethers.parseEther('100'),
           ethers.parseEther('60'),