From 01143fbb041dc81d127bee4d218a32447a556b69 Mon Sep 17 00:00:00 2001 From: Atlas Date: Mon, 30 Dec 2024 16:48:12 +0100 Subject: [PATCH] feat: add withdraw --- src/Pi.sol | 4 ++++ test/Pi.t.sol | 17 +++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/src/Pi.sol b/src/Pi.sol index e00bfc5..b61aa62 100644 --- a/src/Pi.sol +++ b/src/Pi.sol @@ -79,6 +79,10 @@ contract Pi is OVMClient, OwnableUpgradeable { emit ResponseParsed(requestId, success, strPI); } + function withdraw() external onlyOwner { + payable(owner()).transfer(address(this).balance); + } + /** * @dev Retrieves the response associated with the given request ID. * @param requestId The ID of the request. diff --git a/test/Pi.t.sol b/test/Pi.t.sol index ef3e8f2..76fd0cf 100644 --- a/test/Pi.t.sol +++ b/test/Pi.t.sol @@ -27,4 +27,21 @@ contract PiTest is Test { string memory strPI = pi.getResponse("0x1234"); vm.assertEq(strPI, "3.14159"); } + + function testWithdraw() public { + // Fund the contract + vm.deal(address(pi), 1 ether); + + // Ensure only owner can withdraw + vm.expectRevert(); + pi.withdraw(); + + // Check withdrawal by owner + uint256 aliceBalanceBefore = alice.balance; + vm.prank(alice); + pi.withdraw(); + + assertEq(address(pi).balance, 0); + assertEq(alice.balance, aliceBalanceBefore + 1 ether); + } }