Test _callbackResult is cleaned up

alcueca · Jul 30, 2023 · 94a2096 · 94a2096
1 parent 501a9eb
commit 94a2096
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/src/aave/AaveWrapper.sol b/src/aave/AaveWrapper.sol
@@ -78,7 +78,7 @@ contract AaveWrapper is IERC3156PPFlashLender, IFlashLoanSimpleReceiver {
         /// @param data The ABI encoded data to be passed to the callback
         /// @return result ABI encoded result of the callback
         function(address, address, IERC20, uint256, uint256, bytes memory) external returns (bytes memory) callback
-    ) external returns (bytes memory) {
+    ) external returns (bytes memory result) {
         bytes memory data = abi.encode(msg.sender, loanReceiver, callback.encodeFunction(), initiatorData);
 
         POOL.flashLoanSimple({
@@ -89,8 +89,8 @@ contract AaveWrapper is IERC3156PPFlashLender, IFlashLoanSimpleReceiver {
             referralCode: 0
         });
 
-        bytes memory result = _callbackResult;
-        delete _callbackResult; // TODO: Confirm that this deletes the storage variable
+        result = _callbackResult;
+        delete _callbackResult;
         return result;
     }
 

diff --git a/test/AaveWrapper.t.sol b/test/AaveWrapper.t.sol
@@ -59,5 +59,8 @@ contract AaveWrapperTest is PRBTest, StdCheats {
         assertEq(borrower.flashAmount(), loan);
         assertEq(borrower.flashBalance(), loan + fee); // The amount we transferred to pay for fees, plus the amount we borrowed
         assertEq(borrower.flashFee(), fee);
+
+        // Test the wrapper state (return bytes should be cleaned up)
+        assertEq(vm.load(address(wrapper), bytes32(uint256(0))), "");
     }
 }