T

Vectorized · Feb 10, 2024 · 42e5118 · 42e5118
1 parent a867a4f
commit 42e5118
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git 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
@@ -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);
+    	}
     }