fix: retrieve decimals from asset

superical · Oct 1, 2023 · e6809c3 · e6809c3
1 parent 11e9e99
commit e6809c3
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 3 deletions.
diff --git a/contracts/TimeLockVault.sol b/contracts/TimeLockVault.sol
@@ -99,6 +99,19 @@ abstract contract TimeLockVault is ERC20Upgradeable, ITimeLockVault, TimeLockVau
         return _asset.balanceOf(address(this));
     }
 
+    function decimals()
+        public
+        view
+        virtual
+        override(ERC20Upgradeable, IERC20MetadataUpgradeable)
+        returns (uint8)
+    {
+        if (address(_asset) == address(0)) {
+            revert AssetUndefined();
+        }
+        return _asset.decimals();
+    }
+
     function _setAsset(address asset_) internal {
         _asset = IERC20MetadataUpgradeable(asset_);
     }

diff --git a/contracts/interfaces/TimeLockVaultErrors.sol b/contracts/interfaces/TimeLockVaultErrors.sol
@@ -6,6 +6,8 @@ interface TimeLockVaultErrors {
 
     error InvalidActiveDeposit();
 
+    error AssetUndefined();
+
     error DepositNotMatured();
 
     error DepositAlreadyMatured();

diff --git a/contracts/mocks/MockTimeLockVault.sol b/contracts/mocks/MockTimeLockVault.sol
@@ -18,4 +18,8 @@ contract MockTimeLockVault is SimpleTimeLockVault {
     ) public returns (DepositInfo memory) {
         return _prematureWithdraw(from, to, depositId);
     }
+
+    function setAssetInternal(address asset_) public {
+        _setAsset(asset_);
+    }
 }
diff --git a/test/TimeLockVault/time-lock-vault-setup.spec.ts b/test/TimeLockVault/time-lock-vault-setup.spec.ts
@@ -1,5 +1,6 @@
 import { loadFixture } from "@nomicfoundation/hardhat-network-helpers";
 import { expect } from "chai";
+import { constants } from "ethers";
 
 import { MockTimeLockVault } from "../../types";
 import { deployMockTimeLockVaultFixture } from "./time-lock-vault.fixture";
@@ -15,10 +16,37 @@ describe("Time-Lock Vault", () => {
   });
 
   describe("Setup", () => {
-    it("should return the correct asset contract address", async () => {
-      const asset = await mockVaultContract.asset();
+    describe("When asset is specified", () => {
+      it("should return the correct asset contract address", async () => {
+        const asset = await mockVaultContract.asset();
 
-      expect(asset).to.equal(fixtures.mockERC20Contract.address);
+        expect(asset).to.equal(fixtures.mockERC20Contract.address);
+      });
+
+      it("should have the same decimals as asset contract", async () => {
+        const decimals = await mockVaultContract.decimals();
+        const assetDecimals = await fixtures.mockERC20Contract.decimals();
+
+        expect(decimals).to.equal(assetDecimals);
+      });
+    });
+
+    describe("When asset is zero address", () => {
+      beforeEach(async () => {
+        await mockVaultContract
+          .connect(fixtures.signers.deployer)
+          .setAssetInternal(constants.AddressZero);
+
+        // Assert that asset is zero address
+        const asset = await mockVaultContract.asset();
+        expect(asset).to.equal(constants.AddressZero);
+      });
+
+      it("should revert when returning decimals", async () => {
+        const tx = mockVaultContract.decimals();
+
+        await expect(tx).to.be.revertedWithCustomError(mockVaultContract, "AssetUndefined");
+      });
     });
   });
 });