From 1ae3d6fa77212b428ccfed5fda3aaa499d440301 Mon Sep 17 00:00:00 2001
From: pyk <gm@pyk.sh>
Date: Sat, 27 Jan 2024 11:23:03 +0700
Subject: [PATCH] chore: lock & unlock should update totalLockedNFT

---
 src/LlamaLocker.sol    | 11 +++++++++--
 test/LlamaLocker.t.sol | 20 ++++++++++++++++++--
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/LlamaLocker.sol b/src/LlamaLocker.sol
index 1322092..81a1337 100644
--- a/src/LlamaLocker.sol
+++ b/src/LlamaLocker.sol
@@ -37,7 +37,7 @@ contract LlamaLocker is ERC721Holder, Ownable2Step {
     IERC721 public nft;
     mapping(uint256 tokenId => address owner) private nftOwners;
     uint256 public constant REWARD_DURATION = 7 days;
-    uint256 public totalLockedNft;
+    uint256 public totalLockedNFT;
 
     error RenounceInvalid();
     error RewardTokenExists();
@@ -89,7 +89,7 @@ contract LlamaLocker is ERC721Holder, Ownable2Step {
     function distribute(IERC20 _token, uint256 _amount) external onlyOwner {
         if (rewardStates[_token].updatedAt == 0) revert RewardTokenNotExists();
         if (_amount == 0) revert RewardAmountInvalid();
-        if (totalLockedNft == 0) revert NoLockers();
+        if (totalLockedNFT == 0) revert NoLockers();
         //
     }
 
@@ -100,11 +100,18 @@ contract LlamaLocker is ERC721Holder, Ownable2Step {
             nft.safeTransferFrom(msg.sender, address(this), tokenIds_[i]);
             nftOwners[tokenIds_[i]] = msg.sender;
         }
+
+        // Increase total locked NFT
+        totalLockedNFT += tokenCount;
     }
 
     function unlock(uint256[] calldata tokenIds_) external {
         uint256 tokenCount = tokenIds_.length;
         if (tokenCount == 0) revert LockZeroToken();
+
+        // Decrease total locked NFT
+        totalLockedNFT -= tokenCount;
+
         for (uint256 i = 0; i < tokenCount; ++i) {
             if (nftOwners[tokenIds_[i]] != msg.sender) revert UnlockOwnerInvalid();
             nft.safeTransferFrom(address(this), msg.sender, tokenIds_[i]);
diff --git a/test/LlamaLocker.t.sol b/test/LlamaLocker.t.sol
index 2b643bf..28048b3 100644
--- a/test/LlamaLocker.t.sol
+++ b/test/LlamaLocker.t.sol
@@ -139,6 +139,10 @@ contract LlamaLockerTest is Test {
     //   assertEq(data.lastUpdatedAt, block.timestamp);
     // }
 
+    //************************************************************//
+    //                          Lock NFT                          //
+    //************************************************************//
+
     function testLockNFTZeroToken() public {
         uint256[] memory tokenIds = new uint256[](0);
 
@@ -146,7 +150,7 @@ contract LlamaLockerTest is Test {
         locker.lock(tokenIds);
     }
 
-    function testLockNFTTransferred() public {
+    function testLockNFT() public {
         uint256 tokenId1 = nft.mint(alice);
         uint256 tokenId2 = nft.mint(alice);
         uint256[] memory tokenIds = new uint256[](2);
@@ -158,10 +162,18 @@ contract LlamaLockerTest is Test {
         locker.lock(tokenIds);
         vm.stopPrank();
 
+        // NFT should be transfered to locker
         assertEq(nft.ownerOf(tokenId1), address(locker));
         assertEq(nft.ownerOf(tokenId2), address(locker));
+
+        // total locked NFT should be increased
+        assertEq(locker.totalLockedNFT(), 2);
     }
 
+    //************************************************************//
+    //                         Unlock NFT                         //
+    //************************************************************//
+
     function testUnlockNFTInvalidOwner() public {
         uint256 tokenId1 = nft.mint(alice);
         uint256 tokenId2 = nft.mint(alice);
@@ -179,7 +191,7 @@ contract LlamaLockerTest is Test {
         locker.unlock(tokenIds);
     }
 
-    function testUnlockNFTTransfered() public {
+    function testUnlockNFT() public {
         uint256 tokenId1 = nft.mint(alice);
         uint256 tokenId2 = nft.mint(alice);
         uint256[] memory tokenIds = new uint256[](2);
@@ -192,7 +204,11 @@ contract LlamaLockerTest is Test {
         locker.unlock(tokenIds);
         vm.stopPrank();
 
+        // Unlocked NFT should be transferred
         assertEq(nft.ownerOf(tokenId1), alice);
         assertEq(nft.ownerOf(tokenId2), alice);
+
+        // totalLockedNFT should decrease
+        assertEq(locker.totalLockedNFT(), 0);
     }
 }