Added function to allow user control over all mocked price parameters

immunefi-team · Feb 7, 2024 · 7aa244c · 7aa244c
1 parent eed32bf
commit 7aa244c
Showing 1 changed file with 42 additions and 0 deletions.
diff --git a/src/oracle/lib/MockPyth.sol b/src/oracle/lib/MockPyth.sol
@@ -16,6 +16,48 @@ library MockPyth {
         PythUpgradable pythUpgradable;
     }
 
+    /**
+     * @dev Mocks oracle data for testing purposes. This function allows simulating oracle data
+     * within a test environment, enabling the testing of contract interactions with oracles.
+     * @param pid The price feed ID of the oracle data to mock.
+     * @param pythPrice The pyth price { int64 price, uint64 conf, int32 expo, uint256 publishTime }
+     */
+    function mockOracleData(bytes32 pid, PythUpgradable.Price memory pythPrice) internal {
+        Context memory context = context();
+        require(context.pythUpgradable.priceFeedExists(pid), "PythOracle: Price feed not found");
+
+        vm.mockCall(
+            address(context.pythUpgradable),
+            abi.encodePacked(PythUpgradable.getPrice.selector, pid),
+            abi.encode(pythPrice)
+        );
+        vm.mockCall(
+            address(context.pythUpgradable),
+            abi.encodePacked(PythUpgradable.getPriceUnsafe.selector, pid),
+            abi.encode(pythPrice)
+        );
+        vm.mockCall(
+            address(context.pythUpgradable),
+            abi.encodePacked(PythUpgradable.getPriceNoOlderThan.selector, pid),
+            abi.encode(pythPrice)
+        );
+        vm.mockCall(
+            address(context.pythUpgradable),
+            abi.encodePacked(PythUpgradable.getEmaPrice.selector, pid),
+            abi.encode(pythPrice)
+        );
+        vm.mockCall(
+            address(context.pythUpgradable),
+            abi.encodePacked(PythUpgradable.getEmaPriceUnsafe.selector, pid),
+            abi.encode(pythPrice)
+        );
+        vm.mockCall(
+            address(context.pythUpgradable),
+            abi.encodePacked(PythUpgradable.getEmaPriceNoOlderThan.selector, pid),
+            abi.encode(pythPrice)
+        );
+    }
+
     /**
      * @dev Mocks oracle data for testing purposes. This function allows simulating oracle data
      * within a test environment, enabling the testing of contract interactions with oracles.