From 26269db801405dd0cc57e13b3f4cc4f401b3379c Mon Sep 17 00:00:00 2001
From: zeroXbrock <2791467+zeroXbrock@users.noreply.github.com>
Date: Wed, 17 Jul 2024 18:35:53 -0700
Subject: [PATCH] rearrange tests, improve testConfidentialResponse

---
 test/Forge.t.sol | 66 ++++++++++++++++++++++++++++++------------------
 1 file changed, 42 insertions(+), 24 deletions(-)

diff --git a/test/Forge.t.sol b/test/Forge.t.sol
index 26cb9bd..a182806 100644
--- a/test/Forge.t.sol
+++ b/test/Forge.t.sol
@@ -7,30 +7,6 @@ import "src/suavelib/Suave.sol";
 import "src/Context.sol";
 import {Suapp} from "src/Suapp.sol";
 
-contract NumberSuapp is Suapp {
-    uint256 public number;
-
-    function onSetNumber(uint256 num) public {
-        number = num;
-    }
-
-    function setNumber() public confidential returns (bytes memory) {
-        uint256 num = abi.decode(Context.confidentialInputs(), (uint256));
-        return abi.encodeWithSelector(this.onSetNumber.selector, num);
-    }
-}
-
-contract TestConfidential is Test, SuaveEnabled {
-    function testConfidentialResponse() public {
-        NumberSuapp suapp = new NumberSuapp();
-
-        ctx.setConfidentialInputs(abi.encode(123));
-
-        bytes memory suaveCalldata = suapp.setNumber();
-        assertEq(suaveCalldata.length, 4 + 32);
-    }
-}
-
 contract TestForge is Test, SuaveEnabled {
     address[] public addressList = [0xC8df3686b4Afb2BB53e60EAe97EF043FE03Fb829];
 
@@ -104,3 +80,45 @@ contract TestForge is Test, SuaveEnabled {
         assertEq(found3, address(0));
     }
 }
+
+contract TestConfidential is Test, SuaveEnabled {
+    /**
+     * @notice Assumes 36 bytes are given, returns `data[4..]`.
+     */
+    function stripSelector(bytes memory data) internal pure returns (bytes memory trimmedData) {
+        trimmedData = new bytes(data.length - 4);
+        assembly {
+            mstore(add(trimmedData, 0x20), sub(mload(add(data, 0x20)), 0x04))
+            mstore(add(trimmedData, 0x20), mload(add(data, 0x24)))
+        }
+    }
+
+    function testConfidentialResponse() public {
+        NumberSuapp suapp = new NumberSuapp();
+
+        ctx.setConfidentialInputs(abi.encode(123));
+
+        // call confidential/offchain function, verify calldata
+        bytes memory suaveCalldata = suapp.setNumber();
+        assertEq(suaveCalldata.length, 4 + 32);
+        uint256 num = abi.decode(stripSelector(suaveCalldata), (uint256));
+        assertEq(num, 123);
+
+        // call onchain function, verify number
+        suapp.onSetNumber(num);
+        assertEq(suapp.number(), num);
+    }
+}
+
+contract NumberSuapp is Suapp {
+    uint256 public number;
+
+    function onSetNumber(uint256 num) public {
+        number = num;
+    }
+
+    function setNumber() public confidential returns (bytes memory) {
+        uint256 num = abi.decode(Context.confidentialInputs(), (uint256));
+        return abi.encodeWithSelector(this.onSetNumber.selector, num);
+    }
+}