Merge pull request #16 from iotexproject/update_gauge

redefine device nft
iotexproject · Aug 13, 2024 · 11670da · 11670da
2 parents 6bf668a + 357785f
commit 11670da
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 12 deletions.
diff --git a/contracts/gauges/DeviceGauge.sol b/contracts/gauges/DeviceGauge.sol
@@ -2,9 +2,10 @@
 pragma solidity ^0.8.0;
 
 import {ERC721Holder} from "@openzeppelin/contracts/token/ERC721/utils/ERC721Holder.sol";
+import {IERC721} from "@openzeppelin/contracts/token/ERC721/IERC721.sol";
 
 import {IVoter} from "../interfaces/IVoter.sol";
-import {IDeviceNFT} from "../interfaces/IDeviceNFT.sol";
+import {IWeightedNFT} from "../interfaces/IWeightedNFT.sol";
 import {RewardGauge} from "./RewardGauge.sol";
 import {IIncentive} from "../interfaces/IIncentive.sol";
 
@@ -15,12 +16,16 @@ contract DeviceGauge is RewardGauge, ERC721Holder {
   mapping(uint256 => address) public tokenStaker;
   mapping(uint256 => uint256) public tokenWeight;
 
+  address public immutable weightedNFT;
+
   constructor(
     address _forwarder,
-    address _stakingToken,
+    address _weightedNFT,
     address _voter,
     address _incentives
-  ) RewardGauge(_forwarder, _stakingToken, _voter, _incentives) {}
+  ) RewardGauge(_forwarder, IWeightedNFT(_weightedNFT).nft(), _voter, _incentives) {
+    weightedNFT = _weightedNFT;
+  }
 
   function _depositFor(uint256 _tokenId, address _recipient) internal override nonReentrant {
     if (_tokenId == 0) revert ZeroAmount();
@@ -29,8 +34,8 @@ contract DeviceGauge is RewardGauge, ERC721Holder {
     address sender = _msgSender();
     _updateRewards(_recipient);
 
-    IDeviceNFT(stakingToken).safeTransferFrom(sender, address(this), _tokenId);
-    uint256 _amount = IDeviceNFT(stakingToken).weight(_tokenId);
+    IERC721(stakingToken).safeTransferFrom(sender, address(this), _tokenId);
+    uint256 _amount = IWeightedNFT(weightedNFT).weight(_tokenId);
     totalSupply += _amount;
     balanceOf[_recipient] += _amount;
     tokenStaker[_tokenId] = _recipient;
@@ -50,7 +55,7 @@ contract DeviceGauge is RewardGauge, ERC721Holder {
     uint256 _amount = tokenWeight[_tokenId];
     totalSupply -= _amount;
     balanceOf[sender] -= _amount;
-    IDeviceNFT(stakingToken).safeTransferFrom(address(this), sender, _tokenId);
+    IERC721(stakingToken).safeTransferFrom(address(this), sender, _tokenId);
     delete tokenStaker[_tokenId];
     delete tokenWeight[_tokenId];
     updateWeightBalance(sender);

diff --git a/contracts/interfaces/IDeviceNFT.sol → contracts/interfaces/IWeightedNFT.sol b/contracts/interfaces/IDeviceNFT.sol → contracts/interfaces/IWeightedNFT.sol
@@ -1,8 +1,7 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.0;
 
-import {IERC721} from "@openzeppelin/contracts/token/ERC721/IERC721.sol";
-
-interface IDeviceNFT is IERC721 {
+interface IWeightedNFT {
   function weight(uint256 tokenId) external view returns (uint256);
+  function nft() external view returns (address);
 }
diff --git a/contracts/test/TestDeviceNFT.sol b/contracts/test/TestDeviceNFT.sol
@@ -1,9 +1,9 @@
 pragma solidity ^0.8.0;
 
 import {ERC721} from "@openzeppelin/contracts/token/ERC721/ERC721.sol";
-import {IDeviceNFT} from "../interfaces/IDeviceNFT.sol";
+import {IWeightedNFT} from "../interfaces/IWeightedNFT.sol";
 
-contract TestDeviceNFT is IDeviceNFT, ERC721 {
+contract TestDeviceNFT is IWeightedNFT, ERC721 {
   mapping(uint256 => uint256) public weightOf;
 
   constructor(string memory name_, string memory symbol_) ERC721(name_, symbol_) {
@@ -19,6 +19,10 @@ contract TestDeviceNFT is IDeviceNFT, ERC721 {
     return weightOf[tokenId];
   }
 
+  function nft() external view override returns (address) {
+    return address(this);
+  }
+
   function setWeight(uint256 _tokenId, uint256 _weight) public {
     weightOf[_tokenId] = _weight;
   }

diff --git a/test/TestVoter.t.sol b/test/TestVoter.t.sol
@@ -10,6 +10,7 @@ import {IVoter} from "../contracts/interfaces/IVoter.sol";
 import {IRewardGauge} from "../contracts/interfaces/IRewardGauge.sol";
 import {IVault} from "../contracts/interfaces/IVault.sol";
 import {DAOForwarder} from "../contracts/DAOForwarder.sol";
+import {TestDeviceNFT} from "../contracts/test/TestDeviceNFT.sol";
 import {TestStrategyManager} from "../contracts/test/TestStrategyManager.sol";
 import {FactoryRegistry} from "../contracts/factories/FactoryRegistry.sol";
 import {GaugeFactory} from "../contracts/factories/GaugeFactory.sol";
@@ -163,7 +164,7 @@ contract TestVoter is Test {
     voter.createGauge(poolFactory, address(pool), 0, 0);
 
     // 3. create NFT gauge
-    address deviceNFT = address(10);
+    address deviceNFT = address(new TestDeviceNFT("name", "symbol"));
     voter.createGauge(poolFactory, deviceNFT, 1, 10);
     gauge = voter.gauges(address(deviceNFT));
     assertTrue(gauge != address(0));