Total assets tests

solidity/test/MezoAllocator.test.ts

@@ -272,6 +272,11 @@ describe("MezoAllocator", () => {
         const totalAssets = await mezoAllocator.totalAssets()
         expect(totalAssets).to.equal(to1e18(5))
       })
+
+      it("should be equal to the deposit balance", async () => {
+        const depositBalance = await mezoAllocator.depositBalance()
+        expect(await mezoAllocator.totalAssets()).to.equal(depositBalance)
+      })
     })
   })
 

solidity/test/stBTC.test.ts

@@ -1938,6 +1938,152 @@ describe("stBTC", () => {
     })
   })
 
+  describe("totalAssets", () => {
+    beforeAfterSnapshotWrapper()
+
+    const donation = to1e18(1)
+    const firstDeposit = to1e18(2)
+    const secondDeposit = to1e18(3)
+
+    context("when there are no deposits", () => {
+      it("should return 0", async () => {
+        expect(await stbtc.totalAssets()).to.be.eq(0)
+      })
+    })
+
+    context("when there are deposits", () => {
+      context("when there is a first deposit made", () => {
+        beforeAfterSnapshotWrapper()
+
+        before(async () => {
+          await tbtc.mint(depositor1.address, firstDeposit)
+          await tbtc
+            .connect(depositor1)
+            .approve(await stbtc.getAddress(), firstDeposit)
+          await stbtc
+            .connect(depositor1)
+            .deposit(firstDeposit, depositor1.address)
+        })
+
+        it("should return the total assets", async () => {
+          const expectedAssets =
+            firstDeposit - feeOnTotal(firstDeposit, entryFeeBasisPoints)
+          expect(await stbtc.totalAssets()).to.be.eq(expectedAssets)
+        })
+
+        it("should be equal to tBTC balance of the contract", async () => {
+          expect(await stbtc.totalAssets()).to.be.eq(
+            await tbtc.balanceOf(await stbtc.getAddress()),
+          )
+        })
+      })
+
+      context("when there are two deposits made", () => {
+        beforeAfterSnapshotWrapper()
+
+        before(async () => {
+          await tbtc.mint(depositor1.address, firstDeposit + secondDeposit)
+          await tbtc
+            .connect(depositor1)
+            .approve(await stbtc.getAddress(), firstDeposit + secondDeposit)
+          await stbtc
+            .connect(depositor1)
+            .deposit(firstDeposit, depositor1.address)
+          await stbtc
+            .connect(depositor1)
+            .deposit(secondDeposit, depositor1.address)
+        })
+
+        it("should return the total assets", async () => {
+          const expectedAssetsFirstDeposit =
+            firstDeposit - feeOnTotal(firstDeposit, entryFeeBasisPoints)
+          const expectedAssetsSecondDeposit =
+            secondDeposit - feeOnTotal(secondDeposit, entryFeeBasisPoints)
+          expect(await stbtc.totalAssets()).to.be.eq(
+            expectedAssetsFirstDeposit + expectedAssetsSecondDeposit,
+          )
+        })
+      })
+
+      context("when the funds were allocated after deposits", () => {
+        beforeAfterSnapshotWrapper()
+
+        before(async () => {
+          await tbtc.mint(depositor1.address, firstDeposit + secondDeposit)
+          await tbtc
+            .connect(depositor1)
+            .approve(await stbtc.getAddress(), firstDeposit + secondDeposit)
+          await stbtc
+            .connect(depositor1)
+            .deposit(firstDeposit, depositor1.address)
+          await stbtc
+            .connect(depositor1)
+            .deposit(secondDeposit, depositor1.address)
+          await mezoAllocator.connect(maintainer).allocate()
+        })
+
+        it("should return the total assets", async () => {
+          const deposits = firstDeposit + secondDeposit
+          const expectedAssets =
+            deposits - feeOnTotal(deposits, entryFeeBasisPoints)
+          expect(await stbtc.totalAssets()).to.be.eq(expectedAssets)
+        })
+      })
+
+      context("when there is a donation made", () => {
+        beforeAfterSnapshotWrapper()
+
+        let totalAssetsBeforeDonation: bigint
+
+        before(async () => {
+          await tbtc.mint(depositor1.address, firstDeposit)
+          await tbtc
+            .connect(depositor1)
+            .approve(await stbtc.getAddress(), firstDeposit)
+          await stbtc
+            .connect(depositor1)
+            .deposit(firstDeposit, depositor1.address)
+          totalAssetsBeforeDonation = await stbtc.totalAssets()
+          await tbtc.mint(await stbtc.getAddress(), donation)
+        })
+
+        it("should return the total assets", async () => {
+          expect(await stbtc.totalAssets()).to.be.eq(
+            totalAssetsBeforeDonation + donation,
+          )
+        })
+      })
+
+      context("when there was a withdrawal", () => {
+        beforeAfterSnapshotWrapper()
+
+        let totalAssetsBeforeWithdrawal: bigint
+
+        before(async () => {
+          await tbtc.mint(depositor1.address, firstDeposit)
+          await tbtc
+            .connect(depositor1)
+            .approve(await stbtc.getAddress(), firstDeposit)
+          await stbtc
+            .connect(depositor1)
+            .deposit(firstDeposit, depositor1.address)
+          totalAssetsBeforeWithdrawal = await stbtc.totalAssets()
+          await stbtc
+            .connect(depositor1)
+            .withdraw(to1e18(1), depositor1, depositor1)
+        })
+
+        it("should return the total assets", async () => {
+          const actualWithdrawnAssets =
+            to1e18(1) + feeOnRaw(to1e18(1), exitFeeBasisPoints)
+          expect(await stbtc.totalAssets()).to.be.eq(
+            totalAssetsBeforeWithdrawal - actualWithdrawnAssets,
+          )
+        })
+      })
+    })
+  })
+
   describe("feeOnTotal - internal test helper", () => {
     context("when the fee's modulo remainder is greater than 0", () => {
       it("should add 1 to the result", () => {