From a0e356c228a8cbe0815d0ca455f99d777ec8fcd7 Mon Sep 17 00:00:00 2001
From: Kellar <kellar.base09@gmail.com>
Date: Fri, 25 Oct 2024 17:41:14 +0100
Subject: [PATCH] Update tests

---
 contracts/interfaces/hats/IHats.sol |  5 +++
 contracts/mocks/MockHats.sol        | 25 +++++++++++++--
 test/DecentHats_0_1_0.test.ts       | 48 ++++++++++++++++++++++++-----
 3 files changed, 68 insertions(+), 10 deletions(-)

diff --git a/contracts/interfaces/hats/IHats.sol b/contracts/interfaces/hats/IHats.sol
index c460c46d..e78f6c61 100644
--- a/contracts/interfaces/hats/IHats.sol
+++ b/contracts/interfaces/hats/IHats.sol
@@ -39,4 +39,9 @@ interface IHats {
     ) external returns (bool success);
 
     function transferHat(uint256 _hatId, address _from, address _to) external;
+
+    function isWearerOfHat(
+        address _user,
+        uint256 _hatId
+    ) external view returns (bool isWearer);
 }
diff --git a/contracts/mocks/MockHats.sol b/contracts/mocks/MockHats.sol
index f5ac2ece..0eb011a0 100644
--- a/contracts/mocks/MockHats.sol
+++ b/contracts/mocks/MockHats.sol
@@ -5,14 +5,16 @@ import {IHats} from "../interfaces/hats/IHats.sol";
 
 contract MockHats is IHats {
     uint256 public count = 0;
+    mapping(uint256 => address) hatWearers;
 
     function mintTopHat(
-        address,
+        address _target,
         string memory,
         string memory
     ) external returns (uint256 topHatId) {
         topHatId = count;
         count++;
+        hatWearers[topHatId] = _target;
     }
 
     function createHat(
@@ -28,9 +30,26 @@ contract MockHats is IHats {
         count++;
     }
 
-    function mintHat(uint256, address) external pure returns (bool success) {
+    function mintHat(
+        uint256 hatId,
+        address wearer
+    ) external returns (bool success) {
         success = true;
+        hatWearers[hatId] = wearer;
+    }
+
+    function transferHat(uint256 _hatId, address _from, address _to) external {
+        require(
+            hatWearers[_hatId] == _from,
+            "MockHats: Invalid current wearer"
+        );
+        hatWearers[_hatId] = _to;
     }
 
-    function transferHat(uint256, address, address) external {}
+    function isWearerOfHat(
+        address _user,
+        uint256 _hatId
+    ) external view returns (bool isWearer) {
+        isWearer = hatWearers[_hatId] == _user;
+    }
 }
diff --git a/test/DecentHats_0_1_0.test.ts b/test/DecentHats_0_1_0.test.ts
index 6f1ef60d..1285352f 100644
--- a/test/DecentHats_0_1_0.test.ts
+++ b/test/DecentHats_0_1_0.test.ts
@@ -508,7 +508,8 @@ describe('DecentHats_0_1_0', () => {
     });
 
     describe('Creating a new hat on existing Tree', () => {
-      let createHatAndAccountAndMintAndStreamsTx: ethers.ContractTransactionResponse;
+      let createRoleHatPromise: Promise<ethers.ContractTransactionResponse>;
+      const topHatId = 0;
 
       beforeEach(async () => {
         await executeSafeTransaction({
@@ -550,7 +551,7 @@ describe('DecentHats_0_1_0', () => {
 
         const currentBlockTimestamp = (await hre.ethers.provider.getBlock('latest'))!.timestamp;
 
-        createHatAndAccountAndMintAndStreamsTx = await executeSafeTransaction({
+        createRoleHatPromise = executeSafeTransaction({
           safe: gnosisSafe,
           to: decentHatsAddress,
           transactionData: DecentHats_0_1_0__factory.createInterface().encodeFunctionData(
@@ -592,18 +593,51 @@ describe('DecentHats_0_1_0', () => {
         });
       });
 
+      it('Reverts if the top hat is not transferred to the DecentHats module first', async () => {
+        await expect(createRoleHatPromise).to.be.reverted;
+      });
+
       it('Emits an ExecutionSuccess event', async () => {
-        await expect(createHatAndAccountAndMintAndStreamsTx).to.emit(
-          gnosisSafe,
-          'ExecutionSuccess',
-        );
+        // First transfer the top hat to the Safe
+        await mockHats.transferHat(topHatId, gnosisSafeAddress, decentHatsAddress);
+        await expect(await createRoleHatPromise).to.emit(gnosisSafe, 'ExecutionSuccess');
       });
 
       it('Emits an ExecutionFromModuleSuccess event', async () => {
-        await expect(createHatAndAccountAndMintAndStreamsTx)
+        // First transfer the top hat to the Safe
+        await mockHats.transferHat(topHatId, gnosisSafeAddress, decentHatsAddress);
+        await expect(await createRoleHatPromise)
           .to.emit(gnosisSafe, 'ExecutionFromModuleSuccess')
           .withArgs(decentHatsAddress);
       });
+
+      it('Transfers the top hat back to the Safe', async () => {
+        // First transfer the top hat to the Safe
+        await mockHats.transferHat(topHatId, gnosisSafeAddress, decentHatsAddress);
+
+        const isModuleWearerOfTopHat = await mockHats.isWearerOfHat(decentHatsAddress, topHatId);
+        expect(isModuleWearerOfTopHat).to.equal(true);
+
+        await createRoleHatPromise;
+
+        const isSafeWearerOfTopHat = await mockHats.isWearerOfHat(gnosisSafeAddress, topHatId);
+        expect(isSafeWearerOfTopHat).to.equal(true);
+      });
+
+      it('Actually creates the new hat', async () => {
+        const hatsCountBeforeCreate = await mockHats.count();
+        console.log({ hatsCountBeforeCreate });
+
+        // First transfer the top hat to the Safe
+        await mockHats.transferHat(topHatId, gnosisSafeAddress, decentHatsAddress);
+
+        expect(hatsCountBeforeCreate).to.equal(2); // Top hat + admin hat
+
+        await createRoleHatPromise;
+
+        const newHatId = await mockHats.count();
+        expect(newHatId).to.equal(3);
+      });
     });
   });
 });