Merge pull request #6 from ChorusOne/apply_second_review_changes

Apply second review changes
ChorusOne · Nov 28, 2023 · 336d8e2 · 336d8e2
2 parents a7bbdd6 + 231de98
commit 336d8e2
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 21 deletions.
diff --git a/src/FeeRewardsManager.sol b/src/FeeRewardsManager.sol
@@ -6,7 +6,7 @@ import "@openzeppelin/contracts/access/Ownable.sol";
 contract RewardsCollector is Ownable {
     event CollectedReward(
         address withdrawalCredential,
-        uint256 withdrawalFee,
+        uint256 withdrawnAmount,
         address owner,
         uint256 ownerFee
     );
@@ -53,7 +53,7 @@ contract RewardsCollector is Ownable {
         feeNumerator = _feeNumerator;
     }
 
-    function changeFee(uint32 _newFeeNumerator) public onlyOwner {
+    function changeFeeNumerator(uint32 _newFeeNumerator) public onlyOwner {
         feeNumerator = _newFeeNumerator;
     }
 }
@@ -88,6 +88,12 @@ contract FeeRewardsManager is Ownable {
         return payable(addr);
     }
 
+    // Predicts the address of a new contract that will be a `fee_recipient` of
+    // an Ethereum validator.
+    // Given the `_withdrawalCredential` we can instantiate a contract that will
+    // be deployed at a deterministic address, calculated given the
+    // `_withdrawalCredential`, the current contract address and the current
+    // contract's bytecode.
     function predictFeeContractAddress(
         address _withdrawalCredential
     ) public view returns (address) {
@@ -110,11 +116,11 @@ contract FeeRewardsManager is Ownable {
         return address(uint160(uint(hash)));
     }
 
-    function changeFee(
+    function changeFeeNumerator(
         address payable _feeContract,
         uint32 _newFee
     ) public onlyOwner {
-        RewardsCollector(_feeContract).changeFee(_newFee);
+        RewardsCollector(_feeContract).changeFeeNumerator(_newFee);
     }
 
     function batchCollectRewards(

diff --git a/test/FeeRewardsManager.t.sol b/test/FeeRewardsManager.t.sol
@@ -68,6 +68,12 @@ contract FeeRewardsTest is Test {
         // We've got the ether.
         assertEq(address(feeRewardsManager).balance, 0 ether);
         assertEq(address(101).balance, 2.8 ether);
+
+        vm.deal(derivedAddr, 10 ether);
+        RewardsCollector(payable(addr)).collectRewards();
+        assertEq(addr.balance, 0 ether);
+        assertEq(withdrawalCredential.balance, 7.2 ether + 7.2 ether);
+        assertEq(address(101).balance, 2.8 ether);
     }
 
     function createWithdrawalSimulateRewards(
@@ -95,6 +101,10 @@ contract FeeRewardsTest is Test {
         }
         feeRewardsManager.batchCollectRewards(addrs);
         assertEq(address(feeRewardsManager).balance, 280 ether);
+
+        for (uint256 i = 0; i < 100; ++i) {
+            assertEq(address(uint160(i + 100)).balance, 7.2 ether);
+        }
     }
 
     function testChangeDefaultFee() public {
@@ -114,7 +124,7 @@ contract FeeRewardsTest is Test {
         address addr = address(
             createWithdrawalSimulateRewards(address(100), 10 ether)
         );
-        feeRewardsManager.changeFee(payable(addr), 10000);
+        feeRewardsManager.changeFeeNumerator(payable(addr), 10000);
         RewardsCollector(payable(addr)).collectRewards();
         assertEq(address(100).balance, 0 ether);
         // We receive 100%.
@@ -177,28 +187,14 @@ contract FeeRewardsTest is Test {
     }
 
     function testSendToContractWithdrawalCredential() public {
-        WithdrawalContract withdrawalCredentialContract = new WithdrawalContract();
-        address addr = address(
-            createWithdrawalSimulateRewards(
-                address(withdrawalCredentialContract),
-                10 ether
-            )
-        );
-        RewardsCollector(payable(addr)).collectRewards();
-        assertEq(address(feeRewardsManager).balance, 2.8 ether);
-        assertEq(address(withdrawalCredentialContract).balance, 7.2 ether);
-    }
-
-    function testChangeOwnerWithContract() public {
-        WithdrawalContract withdrawalCredentialContract = new WithdrawalContract();
+        ChangeOwnerContract withdrawalCredentialContract = new ChangeOwnerContract();
         address addr = address(
             createWithdrawalSimulateRewards(
                 address(withdrawalCredentialContract),
                 10 ether
             )
         );
+        vm.expectRevert("Failed to send Ether back to withdrawal credential");
         RewardsCollector(payable(addr)).collectRewards();
-        assertEq(address(feeRewardsManager).balance, 2.8 ether);
-        assertEq(address(withdrawalCredentialContract).balance, 7.2 ether);
     }
 }