Adding more tests for Mezo Allocator

core/contracts/MezoAllocator.sol

@@ -217,4 +217,9 @@ contract MezoAllocator is IDispatcher, Ownable2Step {
     function totalAssets() external view returns (uint256 totalAmount) {
         return depositBalance;
     }
+
+    /// @notice Returns the list of maintainers.
+    function getMaintainers() external view returns (address[] memory) {
+        return maintainers;
+    }
 }

core/test/MezoAllocator.test.ts

@@ -3,7 +3,7 @@ import { HardhatEthersSigner } from "@nomicfoundation/hardhat-ethers/signers"
 import { expect } from "chai"
 import { loadFixture } from "@nomicfoundation/hardhat-toolbox/network-helpers"
 
-import { ContractTransactionResponse } from "ethers"
+import { ContractTransactionResponse, ZeroAddress } from "ethers"
 import { beforeAfterSnapshotWrapper, deployment } from "./helpers"
 
 import {
@@ -20,11 +20,12 @@ const { getNamedSigners, getUnnamedSigners } = helpers.signers
 async function fixture() {
   const { tbtc, stbtc, mezoAllocator, mezoPortal } = await deployment()
   const { governance, maintainer } = await getNamedSigners()
-  const [thirdParty] = await getUnnamedSigners()
+  const [depositor, thirdParty] = await getUnnamedSigners()
 
   return {
     governance,
     thirdParty,
+    depositor,
     maintainer,
     tbtc,
     stbtc,
@@ -40,11 +41,21 @@ describe("MezoAllocator", () => {
   let mezoPortal: IMezoPortal
 
   let thirdParty: HardhatEthersSigner
+  let depositor: HardhatEthersSigner
   let maintainer: HardhatEthersSigner
+  let governance: HardhatEthersSigner
 
   before(async () => {
-    ;({ thirdParty, maintainer, tbtc, stbtc, mezoAllocator, mezoPortal } =
-      await loadFixture(fixture))
+    ;({
+      thirdParty,
+      depositor,
+      maintainer,
+      governance,
+      tbtc,
+      stbtc,
+      mezoAllocator,
+      mezoPortal,
+    } = await loadFixture(fixture))
   })
 
   describe("allocate", () => {
@@ -121,6 +132,235 @@ describe("MezoAllocator", () => {
           const depositBalance = await mezoAllocator.depositBalance()
           expect(depositBalance).to.equal(to1e18(11))
         })
+
+        it("should not store any tBTC in Mezo Allocator", async () => {
+          expect(
+            await tbtc.balanceOf(await mezoAllocator.getAddress()),
+          ).to.equal(0)
+        })
+
+        it("should not store any tBTC in stBTC", async () => {
+          expect(await tbtc.balanceOf(await stbtc.getAddress())).to.equal(0)
+        })
+      })
+    })
+  })
+
+  describe("withdraw", () => {
+    beforeAfterSnapshotWrapper()
+
+    context("when a caller is not stBTC", () => {
+      it("should revert", async () => {
+        await expect(
+          mezoAllocator.connect(thirdParty).withdraw(1n),
+        ).to.be.revertedWithCustomError(mezoAllocator, "NotAuthorized")
+      })
+    })
+
+    context("when the caller is stBTC contract", () => {
+      context("when there is no deposit", () => {
+        it("should revert", async () => {
+          await expect(stbtc.withdraw(1n, depositor, depositor))
+            .to.be.revertedWithCustomError(tbtc, "ERC20InsufficientBalance")
+            .withArgs(await mezoPortal.getAddress(), 0, 1n)
+        })
+      })
+
+      context("when there is a deposit", () => {
+        let tx: ContractTransactionResponse
+
+        before(async () => {
+          await tbtc.mint(depositor, to1e18(5))
+          await tbtc.approve(await stbtc.getAddress(), to1e18(5))
+          await stbtc.connect(depositor).deposit(to1e18(5), depositor)
+          await mezoAllocator.connect(maintainer).allocate()
+        })
+
+        context("when the deposit is not fully withdrawn", () => {
+          before(async () => {
+            tx = await stbtc.withdraw(to1e18(2), depositor, depositor)
+          })
+
+          it("should transfer 2 tBTC back to a depositor", async () => {
+            await expect(tx).to.changeTokenBalances(
+              tbtc,
+              [depositor.address],
+              [to1e18(2)],
+            )
+          })
+
+          it("should emit DepositWithdrawn event", async () => {
+            await expect(tx)
+              .to.emit(mezoAllocator, "DepositWithdrawn")
+              .withArgs(1, to1e18(2))
+          })
+
+          it("should decrease tracked deposit balance amount", async () => {
+            const depositBalance = await mezoAllocator.depositBalance()
+            expect(depositBalance).to.equal(to1e18(3))
+          })
+
+          it("should decrease Mezo Portal balance", async () => {
+            expect(
+              await tbtc.balanceOf(await mezoPortal.getAddress()),
+            ).to.equal(to1e18(3))
+          })
+        })
+
+        context("when the deposit is fully withdrawn", () => {
+          before(async () => {
+            tx = await stbtc.withdraw(to1e18(3), depositor, depositor)
+          })
+
+          it("should transfer 3 tBTC back to a depositor", async () => {
+            await expect(tx).to.changeTokenBalances(
+              tbtc,
+              [depositor.address],
+              [to1e18(3)],
+            )
+          })
+
+          it("should emit DepositWithdrawn event", async () => {
+            await expect(tx)
+              .to.emit(mezoAllocator, "DepositWithdrawn")
+              .withArgs(1, to1e18(3))
+          })
+
+          it("should decrease tracked deposit balance amount to zero", async () => {
+            const depositBalance = await mezoAllocator.depositBalance()
+            expect(depositBalance).to.equal(0)
+          })
+
+          it("should decrease Mezo Portal balance", async () => {
+            expect(
+              await tbtc.balanceOf(await mezoPortal.getAddress()),
+            ).to.equal(0)
+          })
+        })
+      })
+    })
+  })
+
+  describe("totalAssets", () => {
+    beforeAfterSnapshotWrapper()
+
+    context("when there is no deposit", () => {
+      it("should return 0", async () => {
+        const totalAssets = await mezoAllocator.totalAssets()
+        expect(totalAssets).to.equal(0)
+      })
+    })
+
+    context("when there is a deposit", () => {
+      before(async () => {
+        await tbtc.mint(await stbtc.getAddress(), to1e18(5))
+        await mezoAllocator.connect(maintainer).allocate()
+      })
+
+      it("should return the total assets value", async () => {
+        const totalAssets = await mezoAllocator.totalAssets()
+        expect(totalAssets).to.equal(to1e18(5))
+      })
+    })
+  })
+
+  describe("addMaintainer", () => {
+    beforeAfterSnapshotWrapper()
+
+    context("when a caller is not a governance", () => {
+      it("should revert", async () => {
+        await expect(
+          mezoAllocator.connect(thirdParty).addMaintainer(depositor.address),
+        ).to.be.revertedWithCustomError(mezoAllocator, "OwnableUnauthorizedAccount")
+      })
+    })
+
+    context("when a caller is governance", () => {
+      context("when a maintainer is added", () => {
+        let tx: ContractTransactionResponse
+
+        before(async () => {
+          tx = await mezoAllocator
+            .connect(governance)
+            .addMaintainer(thirdParty.address)
+        })
+
+        it("should add a maintainer", async () => {
+          expect(await mezoAllocator.isMaintainer(thirdParty.address)).to.equal(
+            true,
+          )
+        })
+
+        it("should emit MaintainerAdded event", async () => {
+          await expect(tx)
+            .to.emit(mezoAllocator, "MaintainerAdded")
+            .withArgs(thirdParty.address)
+        })
+
+        it("should add a new maintainer to the list", async () => {
+          const maintainers = await mezoAllocator.getMaintainers()
+          expect(maintainers).to.deep.equal([maintainer.address, thirdParty.address])
+        })
+
+        it("should not allow to add the same maintainer twice", async () => {
+          await expect(
+            mezoAllocator.connect(governance).addMaintainer(thirdParty.address),
+          ).to.be.revertedWithCustomError(mezoAllocator, "MaintainerAlreadyRegistered")
+        })
+
+        it("should not allow to add a zero address as a maintainer", async () => {
+          await expect(
+            mezoAllocator.connect(governance).addMaintainer(ZeroAddress),
+          ).to.be.revertedWithCustomError(mezoAllocator, "ZeroAddress")
+        })
+      })
+    })
+  })
+
+  describe("removeMaintainer", () => {
+    beforeAfterSnapshotWrapper()
+
+    context("when a caller is not a governance", () => {
+      it("should revert", async () => {
+        await expect(
+          mezoAllocator.connect(thirdParty).removeMaintainer(depositor.address),
+        ).to.be.revertedWithCustomError(mezoAllocator, "OwnableUnauthorizedAccount")
+      })
+    })
+
+    context("when a caller is governance", () => {
+      context("when a maintainer is removed", () => {
+        let tx: ContractTransactionResponse
+
+        before(async () => {
+          await mezoAllocator.connect(governance).addMaintainer(thirdParty.address)
+          tx = await mezoAllocator
+            .connect(governance)
+            .removeMaintainer(thirdParty.address)
+        })
+
+        it("should remove a maintainer", async () => {
+          expect(await mezoAllocator.isMaintainer(thirdParty.address)).to.equal(
+            false,
+          )
+        })
+
+        it("should emit MaintainerRemoved event", async () => {
+          await expect(tx)
+            .to.emit(mezoAllocator, "MaintainerRemoved")
+            .withArgs(thirdParty.address)
+        })
+
+        it("should remove a maintainer from the list", async () => {
+          const maintainers = await mezoAllocator.getMaintainers()
+          expect(maintainers).to.deep.equal([maintainer.address])
+        })
+
+        it("should not allow to remove a maintainer twice", async () => {
+          await expect(
+            mezoAllocator.connect(governance).removeMaintainer(thirdParty.address),
+          ).to.be.revertedWithCustomError(mezoAllocator, "MaintainerNotRegistered")
+        })
       })
     })
   })