adds decimals() override

PufferFinance · Jun 27, 2024 · 6ea6fad · 6ea6fad
1 parent de6833d
commit 6ea6fad
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/mainnet-contracts/src/PufToken.sol b/mainnet-contracts/src/PufToken.sol
@@ -29,7 +29,7 @@ contract PufToken is IPufStakingPool, ERC20, ERC20Permit {
     /**
      * @notice The underlying token decimals
      */
-    uint256 internal immutable _TOKEN_DECIMALS;
+    uint8 internal immutable _TOKEN_DECIMALS;
 
     /**
      * @notice Puffer Token factory
@@ -54,7 +54,7 @@ contract PufToken is IPufStakingPool, ERC20, ERC20Permit {
         // The Factory is the deployer of the contract
         PUFFER_FACTORY = PufferL2Depositor(msg.sender);
         TOKEN = ERC20(token);
-        _TOKEN_DECIMALS = uint256(TOKEN.decimals());
+        _TOKEN_DECIMALS = TOKEN.decimals();
         totalDepositCap = depositCap;
     }
 
@@ -214,4 +214,8 @@ contract PufToken is IPufStakingPool, ERC20, ERC20Permit {
 
         IMigrator(migratorContract).migrate({ depositor: depositor, destination: destination, amount: amount });
     }
+
+    function decimals() public view override returns (uint8 _decimals) {
+        return _TOKEN_DECIMALS;
+    }
 }
diff --git a/mainnet-contracts/test/unit/PufferL2Staking.t.sol b/mainnet-contracts/test/unit/PufferL2Staking.t.sol
@@ -180,6 +180,7 @@ contract PufferL2Staking is UnitTestHelper {
         pufToken.withdraw(bob, amount);
 
         assertEq(sixDecimal.balanceOf(bob), amount, "bob got same amount");
+        assertEq(sixDecimal.decimals(), 6, "decimals matches orginal token");
     }
 
     // Deposit & withdraw 22 decimal token
@@ -209,6 +210,7 @@ contract PufferL2Staking is UnitTestHelper {
         pufToken.withdraw(bob, amount);
 
         assertEq(twentyTwoDecimal.balanceOf(bob), amount, "bob got same amount");
+        assertEq(twentyTwoDecimal.decimals(), 22, "decimals matches orginal token");
     }
 
     // Good Permit signature signature