From 94a20963e5b6aa9555ded5302c8d8e49cf57f0fb Mon Sep 17 00:00:00 2001
From: "ultrasecr.eth" <241804+ultrasecreth@users.noreply.github.com>
Date: Sun, 30 Jul 2023 11:24:09 +0100
Subject: [PATCH] Test _callbackResult is cleaned up

---
 src/aave/AaveWrapper.sol | 6 +++---
 test/AaveWrapper.t.sol   | 3 +++
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/aave/AaveWrapper.sol b/src/aave/AaveWrapper.sol
index 0d8a9a8..4b8870d 100644
--- 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
index 8a5f2d2..d306a39 100644
--- 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))), "");
     }
 }