Implement some simplification in HatsProposalCreationWhitelist

decentdao · Nov 6, 2024 · 436d899 · 436d899
1 parent 93c0a97
commit 436d899
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 31 deletions.
diff --git a/contracts/azorius/strategies/HatsProposalCreationWhitelist.sol b/contracts/azorius/strategies/HatsProposalCreationWhitelist.sol
@@ -11,7 +11,7 @@ abstract contract HatsProposalCreationWhitelist is OwnableUpgradeable {
     IHats public hatsContract;
 
     /** Array to store whitelisted Hat IDs. */
-    uint256[] public whitelistedHatIds;
+    uint256[] private whitelistedHatIds;
 
     error InvalidHatsContract();
     error NoHatsWhitelisted();
@@ -96,24 +96,10 @@ abstract contract HatsProposalCreationWhitelist is OwnableUpgradeable {
     }
 
     /**
-     * Returns the number of whitelisted hats.
-     * @return The number of whitelisted hats
+     * @dev Returns the IDs of all whitelisted Hats.
+     * @return uint256[] memory An array of whitelisted Hat IDs.
      */
-    function getWhitelistedHatsCount() public view returns (uint256) {
-        return whitelistedHatIds.length;
-    }
-
-    /**
-     * Checks if a hat is whitelisted.
-     * @param _hatId The ID of the Hat to check
-     * @return True if the hat is whitelisted, false otherwise
-     */
-    function isHatWhitelisted(uint256 _hatId) public view returns (bool) {
-        for (uint256 i = 0; i < whitelistedHatIds.length; i++) {
-            if (whitelistedHatIds[i] == _hatId) {
-                return true;
-            }
-        }
-        return false;
+    function getWhitelistedHatIds() public view returns (uint256[] memory) {
+        return whitelistedHatIds;
     }
 }
diff --git a/test/HatsProposalCreationWhitelist.test.ts b/test/HatsProposalCreationWhitelist.test.ts
@@ -81,7 +81,9 @@ describe('HatsProposalCreationWhitelist', () => {
     expect(await mockHatsProposalCreationWhitelist.hatsContract()).to.eq(
       await hatsProtocol.getAddress(),
     );
-    expect(await mockHatsProposalCreationWhitelist.isHatWhitelisted(proposerHatId)).to.equal(true);
+    expect(
+      (await mockHatsProposalCreationWhitelist.getWhitelistedHatIds()).includes(proposerHatId),
+    ).to.equal(true);
   });
 
   it('Cannot call setUp function again', async () => {
@@ -144,31 +146,35 @@ describe('HatsProposalCreationWhitelist', () => {
   });
 
   it('Returns correct number of whitelisted hats', async () => {
-    expect(await mockHatsProposalCreationWhitelist.getWhitelistedHatsCount()).to.equal(1);
+    expect((await mockHatsProposalCreationWhitelist.getWhitelistedHatIds()).length).to.equal(1);
 
     await mockHatsProposalCreationWhitelist.connect(owner).whitelistHat(nonProposerHatId);
 
-    expect(await mockHatsProposalCreationWhitelist.getWhitelistedHatsCount()).to.equal(2);
+    expect((await mockHatsProposalCreationWhitelist.getWhitelistedHatIds()).length).to.equal(2);
 
     await mockHatsProposalCreationWhitelist.connect(owner).removeHatFromWhitelist(proposerHatId);
 
-    expect(await mockHatsProposalCreationWhitelist.getWhitelistedHatsCount()).to.equal(1);
+    expect((await mockHatsProposalCreationWhitelist.getWhitelistedHatIds()).length).to.equal(1);
   });
 
   it('Correctly checks if a hat is whitelisted', async () => {
-    expect(await mockHatsProposalCreationWhitelist.isHatWhitelisted(proposerHatId)).to.equal(true);
-    expect(await mockHatsProposalCreationWhitelist.isHatWhitelisted(nonProposerHatId)).to.equal(
-      false,
-    );
+    expect(
+      (await mockHatsProposalCreationWhitelist.getWhitelistedHatIds()).includes(proposerHatId),
+    ).to.equal(true);
+    expect(
+      (await mockHatsProposalCreationWhitelist.getWhitelistedHatIds()).includes(nonProposerHatId),
+    ).to.equal(false);
 
     await mockHatsProposalCreationWhitelist.connect(owner).whitelistHat(nonProposerHatId);
 
-    expect(await mockHatsProposalCreationWhitelist.isHatWhitelisted(nonProposerHatId)).to.equal(
-      true,
-    );
+    expect(
+      (await mockHatsProposalCreationWhitelist.getWhitelistedHatIds()).includes(nonProposerHatId),
+    ).to.equal(true);
 
     await mockHatsProposalCreationWhitelist.connect(owner).removeHatFromWhitelist(proposerHatId);
 
-    expect(await mockHatsProposalCreationWhitelist.isHatWhitelisted(proposerHatId)).to.equal(false);
+    expect(
+      (await mockHatsProposalCreationWhitelist.getWhitelistedHatIds()).includes(proposerHatId),
+    ).to.equal(false);
   });
 });