From b883596a4ec9493c92f4ec4293c59ea49801f292 Mon Sep 17 00:00:00 2001
From: han <han@fbrns.co>
Date: Fri, 17 Jan 2025 00:36:33 +0700
Subject: [PATCH] adjust contract and add partial test

---
 src/NFT.sol    | 54 ++++++++++++++++++++++++++++++++++------
 test/NFT.t.sol | 67 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 113 insertions(+), 8 deletions(-)
 create mode 100644 test/NFT.t.sol

diff --git a/src/NFT.sol b/src/NFT.sol
index 7420773..82249fa 100644
--- a/src/NFT.sol
+++ b/src/NFT.sol
@@ -2,23 +2,61 @@
 pragma solidity 0.8.28;
 
 import {ERC721} from "solmate/tokens/ERC721.sol";
+import {Ownable} from "openzeppelin-contracts/access/Ownable.sol";
 import {Strings} from "openzeppelin-contracts/utils/Strings.sol";
 
-contract NFT is ERC721 {
+
+
+contract NFT is ERC721, Ownable {
+	using Strings for uint256;
+
+	string public baseURI;
 	uint256 public currentTokenId;
+	uint256 public constant TOTAL_SUPPLY = 10_000;
+	uint256 public constant MINT_PRICE = 0.08 ether;
+
+	error MintPriceNotPaid();
+	error MaxSupply();
+	error NonExistentTokenURI();
+	error WithdrawTransfer();
 
 	constructor(
 		string memory _name,
-		string memory _symbol
-	) ERC721(_name, _symbol) {}
+		string memory _symbol,
+		string memory _baseURI
+	) ERC721(_name, _symbol) Ownable(msg.sender) {
+		baseURI = _baseURI;
+	}
 
 	function mintTo(address recipient) public payable returns (uint256) {
-		uint256 newItemId = ++currentTokenId;
-		_safeMint(recipient, newItemId);
-		return newItemId;
+		if (msg.value != MINT_PRICE) {
+			revert MintPriceNotPaid();
+		}
+		uint256 newTokenId = currentTokenId + 1;
+		if (newTokenId > TOTAL_SUPPLY) {
+			revert MaxSupply();
+		}
+		currentTokenId = newTokenId;
+		_safeMint(recipient, newTokenId);
+		return newTokenId;
 	}
 
-	function tokenURI(uint256 id) public view virtual override returns (string memory) {
-		return Strings.toString(id);
+	function tokenURI(uint256 tokenId) public view virtual override returns (string memory) {
+		if (ownerOf(tokenId) == address(0)) {
+			revert NonExistentTokenURI();
+		}
+		return bytes(baseURI).length > 0 ? string(abi.encodePacked(baseURI, tokenId.toString())) : "";
+	}
+
+	function withdrawPayments(address payable payee) external onlyOwner {
+		if (address(this).balance == 0) {
+			revert WithdrawTransfer();
+		}
+
+		payable(payee).transfer(address(this).balance);
+	}
+
+	function _checkOwner() internal view override {
+		require(msg.sender == owner(), "Onable: caller is not the owner");
 	}
 }
diff --git a/test/NFT.t.sol b/test/NFT.t.sol
new file mode 100644
index 0000000..f22d678
--- /dev/null
+++ b/test/NFT.t.sol
@@ -0,0 +1,67 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity 0.8.28;
+
+import {stdStorage, StdStorage, Test} from "forge-std/Test.sol";
+import {NFT} from "src/NFT.sol";
+
+contract NFTTest is Test {
+	using stdStorage for StdStorage;
+
+	NFT private nft;
+
+	function setUp() public {
+		nft = new NFT("NFT_Tutorial", "TUT", "baseUri");
+	}
+
+	function test_RevertMintWithoutValue() public {
+		vm.expectRevert(NFT.MintPriceNotPaid.selector);
+		nft.mintTo(address(1));
+	}
+
+	function test_MintPricePaid() public {
+		nft.mintTo{value: 0.08 ether}(address(1));
+	}
+
+	function test_RevertMintMaxSupplyReached() public {
+		string memory signature = "currentTokenId()";
+		uint256 slot = stdstore.target(address(nft)).sig(signature).find();
+		bytes32 loc = bytes32(slot);
+		bytes32 mockedCurrentTokenId = bytes32(abi.encode(10000));
+		vm.store(address(nft), loc, mockedCurrentTokenId);
+		vm.expectRevert(NFT.MaxSupply.selector);
+		nft.mintTo{value: 0.08 ether}(address(1));
+	}
+
+	function test_RevertMintToZeroAddress() public {
+		vm.expectRevert("INVALID_RECIPIENT");
+		nft.mintTo{value: 0.08 ether}(address(0));
+	}
+
+	function test_NewMintOwnerRegistered() public {
+		nft.mintTo{value: 0.08 ether}(address(1));
+		uint256 slotOfNewOwner = stdstore.target(address(nft)).sig(nft.ownerOf.selector).with_key(address(1)).find();
+
+		uint160 ownerOfTokenIdOne = uint160(
+			uint256(
+				(vm.load(address(nft), bytes32(abi.encode(slotOfNewOwner))))
+			)
+		);
+		assertEq(address(ownerOfTokenIdOne), address(1));
+	}
+
+	function test_BalanceIncremented() public {
+		nft.mintTo{value: 0.08 ether}(address(1));
+		uint256 slotBalance = stdstore.target(address(nft)).sig(nft.balanceOf.selector).with_key(address(1)).find();
+
+		uint256 balanceFirstMint = uint256(
+			vm.load(address(nft), bytes32(slotBalance))
+		);
+		assertEq(balanceFirstMint, 1);
+
+		nft.mintTo{value: 0.08 ether}(address(1));
+		uint256 balanceSecondMint = uint256(
+			vm.load(address(nft), bytes32(slotBalance))
+		);
+		assertEq(balanceSecondMint, 2);
+	}
+}