From dd14267725cccf8dcbd7a73196429b301e4149e8 Mon Sep 17 00:00:00 2001
From: audsssy <stamfordh@hotmail.com>
Date: Mon, 9 Oct 2023 20:57:10 -0700
Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=A7=20moar=20contract=20and=20test=20u?=
 =?UTF-8?q?pdates?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/KaliBerger.sol    |  8 ++++----
 test/KaliBerger.t.sol | 36 ++++++++++++++++++++++++++++++++----
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/src/KaliBerger.sol b/src/KaliBerger.sol
index d16e518..423304e 100644
--- a/src/KaliBerger.sol
+++ b/src/KaliBerger.sol
@@ -322,10 +322,10 @@ contract KaliBerger is Storage {
         onlyOwner(token, tokenId)
     {
         uint256 deposit = this.getDeposit(token, tokenId);
-        uint256 diff = (amount - deposit > 0) ? type(uint256).max : deposit - amount;
+        uint256 diff = (deposit >= amount) ? deposit - amount : type(uint256).max;
 
         if (diff != type(uint256).max) {
-            (bool success,) = msg.sender.call{value: diff}("");
+            (bool success,) = msg.sender.call{value: amount}("");
             if (!success) revert TransferFailed();
 
             _forecloseIfNecessary(token, tokenId, diff, amount);
@@ -589,8 +589,8 @@ contract KaliBerger is Storage {
     //     }
     // }
 
-    function _forecloseIfNecessary(address token, uint256 tokenId, uint256 deposit, uint256 amount) internal {
-        if (deposit == 0 && amount == 0) {
+    function _forecloseIfNecessary(address token, uint256 tokenId, uint256 diff, uint256 amount) internal {
+        if (diff == 0 && amount == 0) {
             IERC721(token).safeTransferFrom(IERC721(token).ownerOf(tokenId), address(this), tokenId);
             deleteDeposit(token, tokenId);
             deletePrice(token, tokenId);
diff --git a/test/KaliBerger.t.sol b/test/KaliBerger.t.sol
index a0ed6a7..42d2aa2 100644
--- a/test/KaliBerger.t.sol
+++ b/test/KaliBerger.t.sol
@@ -289,22 +289,50 @@ contract KaliBergerTest is Test {
         vm.warp(4500);
 
         uint256 _deposit = kaliBerger.getDeposit(address(token_1), 1);
+        // emit log_uint(_deposit);
+        uint256 patronage = kaliBerger.patronageToCollect(address(token_1), 1);
+        // emit log_uint(patronage);
 
         // Bob exits a portion of deposit.
         vm.prank(bob);
         kaliBerger.exit(address(token_1), 1, 0.3 ether);
 
         // Validate deposit amount
-        assertEq(
-            kaliBerger.getDeposit(address(token_1), 1),
-            _deposit - 0.3 ether - kaliBerger.patronageToCollect(address(token_1), 1)
-        );
+        assertEq(kaliBerger.getDeposit(address(token_1), 1), _deposit - 0.3 ether - patronage);
         validatePatronageToCollect(token_1, 1);
     } // timestamp: 4500
 
     /// @notice Bob ragequits by removing all of his deposit, triggering foreclosure.
     function testBuy_ragequit() public payable {
         testBuy_addDeposit();
+
+        vm.warp(5000);
+
+        uint256 _deposit = kaliBerger.getDeposit(address(token_1), 1);
+        // emit log_uint(_deposit);
+        uint256 patronage = kaliBerger.patronageToCollect(address(token_1), 1);
+        // emit log_uint(patronage);
+
+        // Bob withdraws all of deposit.
+        vm.prank(bob);
+        kaliBerger.exit(address(token_1), 1, _deposit - patronage);
+
+        // Validate
+        assertEq(kaliBerger.getDeposit(address(token_1), 1), 0);
+        assertEq(token_1.balanceOf(address(kaliBerger)), 1);
+        validatePatronageToCollect(token_1, 1);
+    }
+
+    /// @notice Bob withdraws too much and triggers InvalidExit() error.
+    function testBuy_invalidExit() public payable {
+        testBuy_addDeposit();
+
+        vm.warp(5000);
+
+        // InvalidExit()
+        vm.expectRevert(KaliBerger.InvalidExit.selector);
+        vm.prank(bob);
+        kaliBerger.exit(address(token_1), 1, 1 ether);
     }
 
     /// @notice Charlie buys token_1, tokenId #1 and announces new price for sale.