solace-bg-dubai/contracts/test/TreasuryWallet.test.ts

import { expect } from "chai";
import { ethers } from "hardhat";
import { TreasuryWallet } from "../typechain-types";
import { SignerWithAddress } from "@nomicfoundation/hardhat-ethers/signers";

describe("TreasuryWallet", function () {
  let treasury: TreasuryWallet;
  let owner1: SignerWithAddress;
  let owner2: SignerWithAddress;
  let owner3: SignerWithAddress;
  let recipient: SignerWithAddress;
  let threshold: number;

  beforeEach(async function () {
    [owner1, owner2, owner3, recipient] = await ethers.getSigners();
    threshold = 2; // 2-of-3 multisig

    const TreasuryWalletFactory = await ethers.getContractFactory("TreasuryWallet");
    treasury = await TreasuryWalletFactory.deploy(
      [owner1.address, owner2.address, owner3.address],
      threshold
    );
    await treasury.waitForDeployment();
  });

  describe("Deployment", function () {
    it("Should set the correct owners and threshold", async function () {
      expect(await treasury.getOwnerCount()).to.equal(3);
      expect(await treasury.threshold()).to.equal(2);
      expect(await treasury.isOwner(owner1.address)).to.be.true;
      expect(await treasury.isOwner(owner2.address)).to.be.true;
      expect(await treasury.isOwner(owner3.address)).to.be.true;
    });

    it("Should reject invalid threshold", async function () {
      const TreasuryWalletFactory = await ethers.getContractFactory("TreasuryWallet");
      await expect(
        TreasuryWalletFactory.deploy([owner1.address], 2)
      ).to.be.revertedWith("TreasuryWallet: invalid threshold");
    });
  });

  describe("Native Token Transfers", function () {
    it("Should allow receiving ETH", async function () {
      await owner1.sendTransaction({
        to: await treasury.getAddress(),
        value: ethers.parseEther("1.0"),
      });
      expect(await ethers.provider.getBalance(await treasury.getAddress())).to.equal(
        ethers.parseEther("1.0")
      );
    });

    it("Should propose and execute a transaction with sufficient approvals", async function () {
      // Send ETH to treasury
      await owner1.sendTransaction({
        to: await treasury.getAddress(),
        value: ethers.parseEther("1.0"),
      });

      // Propose transaction (auto-approves by proposer)
      const tx = await treasury
        .connect(owner1)
        .proposeTransaction(recipient.address, ethers.parseEther("0.5"), "0x");
      const receipt = await tx.wait();
      const proposalId = 0;

      // Owner2 approves
      await treasury.connect(owner2).approveTransaction(proposalId);

      // Check approval count
      const transaction = await treasury.getTransaction(proposalId);
      expect(transaction.approvalCount).to.equal(2);

      // Execute transaction
      await treasury.connect(owner1).executeTransaction(proposalId);

      // Check recipient received funds
      expect(await ethers.provider.getBalance(recipient.address)).to.be.gt(0);
    });

    it("Should not execute transaction without sufficient approvals", async function () {
      await owner1.sendTransaction({
        to: await treasury.getAddress(),
        value: ethers.parseEther("1.0"),
      });

      await treasury
        .connect(owner1)
        .proposeTransaction(recipient.address, ethers.parseEther("0.5"), "0x");

      // Try to execute without second approval
      await expect(treasury.connect(owner1).executeTransaction(0)).to.be.revertedWith(
        "TreasuryWallet: insufficient approvals"
      );
    });
  });

  describe("Owner Management", function () {
    it("Should add a new owner", async function () {
      const newOwner = (await ethers.getSigners())[4];
      await treasury.connect(owner1).addOwner(newOwner.address);
      expect(await treasury.isOwner(newOwner.address)).to.be.true;
      expect(await treasury.getOwnerCount()).to.equal(4);
    });

    it("Should remove an owner", async function () {
      await treasury.connect(owner1).removeOwner(owner3.address);
      expect(await treasury.isOwner(owner3.address)).to.be.false;
      expect(await treasury.getOwnerCount()).to.equal(2);
    });

    it("Should not allow removing owner if it breaks threshold", async function () {
      // Try to remove owner when only 2 remain and threshold is 2
      await treasury.connect(owner1).removeOwner(owner3.address);
      await expect(
        treasury.connect(owner1).removeOwner(owner2.address)
      ).to.be.revertedWith("TreasuryWallet: cannot remove owner, would break threshold");
    });

    it("Should change threshold", async function () {
      await treasury.connect(owner1).changeThreshold(3);
      expect(await treasury.threshold()).to.equal(3);
    });
  });

  describe("Access Control", function () {
    it("Should not allow non-owner to propose transaction", async function () {
      await expect(
        treasury.connect(recipient).proposeTransaction(recipient.address, 0, "0x")
      ).to.be.revertedWith("TreasuryWallet: not an owner");
    });

    it("Should not allow non-owner to approve transaction", async function () {
      await owner1.sendTransaction({
        to: await treasury.getAddress(),
        value: ethers.parseEther("1.0"),
      });

      await treasury
        .connect(owner1)
        .proposeTransaction(recipient.address, ethers.parseEther("0.5"), "0x");

      await expect(treasury.connect(recipient).approveTransaction(0)).to.be.revertedWith(
        "TreasuryWallet: not an owner"
      );
    });
  });
});