diff --git a/src/DN404.sol b/src/DN404.sol index 276a3aa..6dea4db 100644 --- a/src/DN404.sol +++ b/src/DN404.sol @@ -150,6 +150,10 @@ abstract contract DN404 { } } + function balanceOf(address owner) public view virtual returns (uint256) { + return _getDN404Storage().addressData[owner].balance; + } + function ownerOf(uint256 id) public view virtual returns (address owner) { DN404Storage storage $ = _getDN404Storage(); owner = $.aliasToAddress[$.ownerships.get(id)]; diff --git a/test/DN404.t.sol b/test/DN404.t.sol index 53f5422..27fc6a0 100644 --- a/test/DN404.t.sol +++ b/test/DN404.t.sol @@ -9,11 +9,30 @@ contract DN404Test is SoladyTest { function setUp() public { dn = new MockDN404(); - } - function testRegisterAndResolveAlias() public { + function testRegisterAndResolveAlias(address a0, address a1) public { + assertEq(dn.registerAndResolveAlias(a0), 1); + if (a1 == a0) { + assertEq(dn.registerAndResolveAlias(a1), 1); + } else { + assertEq(dn.registerAndResolveAlias(a1), 2); + assertEq(dn.registerAndResolveAlias(a0), 1); + } + } + function testInitialize(uint32 totalNFTSupply, address initialSupplyOwner) public { + if (totalNFTSupply == 0 || uint256(totalNFTSupply) + 1 > type(uint32).max) { + vm.expectRevert(DN404.InvalidTotalNFTSupply.selector); + dn.initializeDN404(totalNFTSupply, initialSupplyOwner); + } else if (initialSupplyOwner == address(0)) { + vm.expectRevert(DN404.TransferToZeroAddress.selector); + dn.initializeDN404(totalNFTSupply, initialSupplyOwner); + } else { + dn.initializeDN404(totalNFTSupply, initialSupplyOwner); + assertEq(dn.totalSupply(), uint256(totalNFTSupply) * 10 ** 18); + assertEq(dn.balanceOf(initialSupplyOwner), uint256(totalNFTSupply) * 10 ** 18); + } }