new: add unionLens test.

unioncredit · Oct 12, 2023 · 7399a0d · 7399a0d
1 parent cab3ec3
commit 7399a0d
Show file tree

Hide file tree

Showing 2 changed files with 84 additions and 3 deletions.
diff --git a/contracts/mocks/UserManagerMock.sol b/contracts/mocks/UserManagerMock.sol
@@ -1,11 +1,16 @@
 //SPDX-License-Identifier: MIT
 pragma solidity 0.8.16;
 
+import {SafeCastUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/math/SafeCastUpgradeable.sol";
+
 /**
  * @title UserManager Contract
  * @dev Manages the Union members credit lines, and their vouchees and borrowers info.
  */
 contract UserManagerMock {
+    using SafeCastUpgradeable for uint256;
+    using SafeCastUpgradeable for uint128;
+
     uint256 public constant MAX_TRUST_LIMIT = 100;
     uint256 public constant MAX_STAKE_AMOUNT = 1000e18;
 
@@ -19,8 +24,27 @@ contract UserManagerMock {
     uint256 public totalFrozenAmount;
     address public stakingToken;
 
+    struct Index {
+        bool isSet;
+        uint128 idx;
+    }
+    struct Vouch {
+        address staker;
+        uint96 trust;
+        uint96 locked;
+        uint64 lastUpdated;
+    }
+    struct Vouchee {
+        address borrower;
+        uint96 voucherIndex;
+    }
+
     mapping(address => uint256) public balances;
     mapping(address => mapping(address => uint256)) public trust;
+    mapping(address => mapping(address => Index)) public voucherIndexes;
+    mapping(address => mapping(address => Index)) public voucheeIndexes;
+    mapping(address => Vouch[]) public vouchers;
+    mapping(address => Vouchee[]) public vouchees;
 
     constructor(address _stakingToken) {
         stakingToken = _stakingToken;
@@ -75,8 +99,21 @@ contract UserManagerMock {
 
     function addMember(address account) public {}
 
-    function updateTrust(address borrower_, uint96 trustAmount) external {
-        trust[msg.sender][borrower_] = trustAmount;
+    function updateTrust(address borrower, uint96 trustAmount) external {
+        address staker = msg.sender;
+        trust[staker][borrower] = trustAmount;
+        Index memory index = voucherIndexes[borrower][staker];
+        if (index.isSet) {
+            Vouch storage vouch = vouchers[borrower][index.idx];
+            vouch.trust = trustAmount;
+        } else {
+            uint256 voucheeIndex = vouchees[staker].length;
+            uint256 voucherIndex = vouchers[borrower].length;
+            voucherIndexes[borrower][staker] = Index(true, voucherIndex.toUint128());
+            voucheeIndexes[borrower][staker] = Index(true, voucheeIndex.toUint128());
+            vouchers[borrower].push(Vouch(staker, trustAmount, 0, 0));
+            vouchees[staker].push(Vouchee(borrower, voucherIndex.toUint96()));
+        }
     }
 
     function cancelVouch(address staker, address borrower) external {}
@@ -111,7 +148,13 @@ contract UserManagerMock {
 
     function onRepayBorrow(address borrower, uint256 pastBlocks) public {}
 
-    function getVoucherCount(address borrower) external view returns (uint256) {}
+    function getVoucherCount(address borrower) external view returns (uint256) {
+        return vouchers[borrower].length;
+    }
+
+    function getVoucheeCount(address borrower) external view returns (uint256) {
+        return vouchees[borrower].length;
+    }
 
     function setEffectiveCount(uint256 effectiveCount) external {}
 

diff --git a/test/foundry/UnionLens.t.sol b/test/foundry/UnionLens.t.sol
@@ -0,0 +1,38 @@
+pragma solidity ^0.8.0;
+
+import {TestWrapper} from "./TestWrapper.sol";
+import {UnionLens} from "union-v2-contracts/UnionLens.sol";
+
+contract TestUnionLens is TestWrapper {
+    UnionLens public unionLens;
+    address public constant ADMIN = address(0);
+
+    function setUp() public virtual {
+        deployMocks();
+        vm.startPrank(ADMIN);
+        marketRegistryMock.setUToken(address(daiMock), address(uTokenMock));
+        marketRegistryMock.setUserManager(address(daiMock), address(userManagerMock));
+        vm.stopPrank();
+        unionLens = new UnionLens(marketRegistryMock);
+    }
+
+    function testGetStakerAddresses() public {
+        address[] memory addresses = unionLens.getStakerAddresses(address(daiMock), address(1));
+        assertEq(addresses.length, 0);
+
+        userManagerMock.updateTrust(address(1), 1 ether);
+        addresses = unionLens.getStakerAddresses(address(daiMock), address(1));
+        assertEq(addresses.length, 1);
+    }
+
+    function testGetBorrowerAddresses() public {
+        address[] memory addresses = unionLens.getBorrowerAddresses(address(daiMock), address(1));
+        assertEq(addresses.length, 0);
+
+        vm.startPrank(address(1));
+        userManagerMock.updateTrust(address(2), 1 ether);
+        addresses = unionLens.getBorrowerAddresses(address(daiMock), address(1));
+        assertEq(addresses.length, 1);
+        vm.stopPrank();
+    }
+}