Merge pull request #2 from GenerationSoftware/gen-1183-improve-rngwit…

…net-so-that-caller-can-use-their-existing Altered interface to allow requestors to use their balance
GenerationSoftware · Mar 1, 2024 · e864f1e · e864f1e
2 parents dcefb3f + 5889fa2
commit e864f1e
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 21 deletions.
diff --git a/src/Requestor.sol b/src/Requestor.sol
@@ -12,8 +12,8 @@ contract Requestor {
         creator = msg.sender;
     }
 
-    function randomize(IWitnetRandomness _witnetRandomness) external payable onlyCreator {
-        _witnetRandomness.randomize{ value: msg.value }();
+    function randomize(uint value, IWitnetRandomness _witnetRandomness) external payable onlyCreator {
+        _witnetRandomness.randomize{ value: value }();
     }
 
     function withdraw(address payable _to) external onlyCreator {

diff --git a/src/RngWitnet.sol b/src/RngWitnet.sol
@@ -52,25 +52,14 @@ contract RngWitnet is IRng {
         return witnetRandomness.estimateRandomizeFee(_gasPrice);
     }
 
-    /**
-    * @notice Sends a request for a random number to the 3rd-party service
-    * @dev Some services will complete the request immediately, others may have a time-delay
-    * @dev Some services require payment in the form of a token, such as $LINK for Chainlink VRF
-    * @return requestId The ID of the request used to get the results of the RNG service
-    * @return lockBlock The block number at which the RNG service will start generating time-delayed randomness.
-    * The calling contract should "lock" all activity until the result is available via the `requestId`
-    */
-    function requestRandomNumber() public payable returns (uint32 requestId, uint256 lockBlock) {
-        if (msg.value == 0) {
-            revert NoPayment();
-        }
+    function requestRandomNumber(uint256 rngPaymentAmount) public payable returns (uint32 requestId, uint256 lockBlock) {
         Requestor requestor = getRequestor(msg.sender);
         unchecked {
             requestId = ++lastRequestId;
             lockBlock = block.number;
         }
         requests[requestId] = lockBlock;
-        requestor.randomize{value: msg.value}(witnetRandomness);
+        requestor.randomize{value: msg.value}(rngPaymentAmount, witnetRandomness);
 
         emit RandomNumberRequested(requestId, msg.sender);
     }
@@ -99,8 +88,8 @@ contract RngWitnet is IRng {
         return uint256(witnetRandomness.getRandomnessAfter(requests[requestId]));
     }
 
-    function startDraw(DrawManager _drawManager, address _rewardRecipient) external payable {
-        (uint32 requestId,) = requestRandomNumber();
+    function startDraw(uint256 rngPaymentAmount, DrawManager _drawManager, address _rewardRecipient) external payable {
+        (uint32 requestId,) = requestRandomNumber(rngPaymentAmount);
         _drawManager.startDraw(_rewardRecipient, requestId);
     }
 }
diff --git a/test/RngWitnet.t.sol b/test/RngWitnet.t.sol
@@ -21,7 +21,7 @@ contract RngWitnetTest is Test {
 
     function testRequestRandomNumber() external {
         vm.mockCall(address(witnetRandomness), 1e18, abi.encodeWithSelector(IWitnetRandomness.randomize.selector), abi.encode(0.5e18));
-        (uint32 requestId, uint256 lockBlock) = rngWitnet.requestRandomNumber{value: 1e18}();
+        (uint32 requestId, uint256 lockBlock) = rngWitnet.requestRandomNumber{value: 1e18}(1e18);
         assertEq(requestId, 1);
         assertEq(lockBlock, block.number);
         assertEq(address(rngWitnet.getRequestor(address(this))).balance, 1e18, "witnet balance should be 1e18");
@@ -31,13 +31,13 @@ contract RngWitnetTest is Test {
         DrawManager drawManager = DrawManager(makeAddr("DrawManager"));
         vm.mockCall(address(drawManager), abi.encodeWithSelector(drawManager.startDraw.selector, address(this), 1), abi.encode());
         vm.mockCall(address(witnetRandomness), 1e18, abi.encodeWithSelector(IWitnetRandomness.randomize.selector), abi.encode(0.5e18));
-        rngWitnet.startDraw{value: 1e18}(drawManager, address(this));
+        rngWitnet.startDraw{value: 1e18}(1e18, drawManager, address(this));
         assertEq(address(rngWitnet.getRequestor(address(this))).balance, 1e18, "witnet balance should be 1e18");
     }
 
     function testWithdraw() public {
         vm.mockCall(address(witnetRandomness), 1e18, abi.encodeWithSelector(IWitnetRandomness.randomize.selector), abi.encode(0.5e18));
-        rngWitnet.requestRandomNumber{value: 1e18}();
+        rngWitnet.requestRandomNumber{value: 1e18}(1e18);
         rngWitnet.withdraw();
         assertEq(address(this).balance, 1000e18, "balance is restored");
     }

diff --git a/test/fork/RngWitnetFork.t.sol b/test/fork/RngWitnetFork.t.sol
@@ -23,7 +23,7 @@ contract RngWitnetForkTest is Test {
     }
 
     function testRequestRandomNumberFromFork() external {
-        (uint32 requestId, uint256 lockBlock) = rngWitnet.requestRandomNumber{value: 1e18}();
+        (uint32 requestId, uint256 lockBlock) = rngWitnet.requestRandomNumber{value: 1e18}(1e18);
         assertEq(requestId, 1);
         assertEq(lockBlock, block.number);
     }