Skip to content

Commit

Permalink
change implementation of StETH to almost real
Browse files Browse the repository at this point in the history
  • Loading branch information
@RuslanProgrammer
RuslanProgrammer committed Jan 12, 2024
1 parent 6db4a90 commit 07d8c1c
Show file tree
Hide file tree
Showing 2 changed files with 85 additions and 21 deletions.
91 changes: 78 additions & 13 deletions contracts/mock/tokens/StETHMock.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,30 +2,95 @@
pragma solidity ^0.8.20;

import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import {PRECISION} from "@solarity/solidity-lib/utils/Globals.sol";
import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";

contract StETHMock is ERC20, Ownable {
uint256 public balanceMultiplicator = PRECISION;
uint256 public totalShares;
uint256 public totalPooledEther;

constructor() ERC20("Staked Ether Mock", "stETHMock") {}
mapping(address => uint256) private shares;

function mint(address account_, uint256 amount_) external {
require(amount_ <= 1000 * (10 ** decimals()), "StETHMock: amount is too big");
constructor() ERC20("Staked Ether Mock", "stETHMock") {
_mintShares(address(this), 10 ** decimals());

_mint(account_, amount_);
totalPooledEther = 10 ** decimals();
}

function _transfer(address sender_, address recipient_, uint256 amount_) internal override {
amount_ = (amount_ * PRECISION) / balanceMultiplicator;
super._transfer(sender_, recipient_, amount_);
function mint(address _account, uint256 _amount) external {
require(_amount <= 1000 * (10 ** decimals()), "StETHMock: amount is too big");

uint256 sharesAmount = getSharesByPooledEth(_amount);

_mintShares(_account, sharesAmount);

totalPooledEther += _amount;
}

function setTotalPooledEther(uint256 _totalPooledEther) external onlyOwner {
totalPooledEther = _totalPooledEther;
}

function totalSupply() public view override returns (uint256) {
return totalPooledEther;
}

function balanceOf(address _account) public view override returns (uint256) {
return getPooledEthByShares(_sharesOf(_account));
}

function sharesOf(address _account) external view returns (uint256) {
return _sharesOf(_account);
}

function getSharesByPooledEth(uint256 _ethAmount) public view returns (uint256) {
return (_ethAmount * totalShares) / totalPooledEther;
}

function getPooledEthByShares(uint256 _sharesAmount) public view returns (uint256) {
return (_sharesAmount * totalPooledEther) / totalShares;
}

function transferShares(address _recipient, uint256 _sharesAmount) external returns (uint256) {
_transferShares(msg.sender, _recipient, _sharesAmount);
uint256 tokensAmount = getPooledEthByShares(_sharesAmount);
return tokensAmount;
}

function transferSharesFrom(address _sender, address _recipient, uint256 _sharesAmount) external returns (uint256) {
uint256 tokensAmount = getPooledEthByShares(_sharesAmount);
_spendAllowance(_sender, msg.sender, tokensAmount);
_transferShares(_sender, _recipient, _sharesAmount);
return tokensAmount;
}

function _transfer(address _sender, address _recipient, uint256 _amount) internal override {
uint256 _sharesToTransfer = getSharesByPooledEth(_amount);
_transferShares(_sender, _recipient, _sharesToTransfer);
}

function setBalanceMultiplicator(uint256 balanceMultiplicator_) external onlyOwner {
balanceMultiplicator = balanceMultiplicator_;
function _sharesOf(address _account) internal view returns (uint256) {
return shares[_account];
}

function balanceOf(address account_) public view override returns (uint256) {
return (super.balanceOf(account_) * balanceMultiplicator) / PRECISION;
function _transferShares(address _sender, address _recipient, uint256 _sharesAmount) internal {
require(_sender != address(0), "TRANSFER_FROM_ZERO_ADDR");
require(_recipient != address(0), "TRANSFER_TO_ZERO_ADDR");
require(_recipient != address(this), "TRANSFER_TO_STETH_CONTRACT");

uint256 currentSenderShares = shares[_sender];
require(_sharesAmount <= currentSenderShares, "BALANCE_EXCEEDED");

shares[_sender] = currentSenderShares - _sharesAmount;
shares[_recipient] += _sharesAmount;
}

function _mintShares(address _recipient, uint256 _sharesAmount) internal returns (uint256 newTotalShares) {
require(_recipient != address(0), "MINT_TO_ZERO_ADDR");

totalShares += _sharesAmount;

shares[_recipient] += _sharesAmount;

return totalShares;
}
}
15 changes: 7 additions & 8 deletions test/Distribution.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,6 @@ describe('Distribution', () => {
depositToken.approve(distribution, wei(1000)),
depositToken.connect(SECOND).approve(distribution, wei(1000)),
]);

await l1Sender.transferOwnership(distribution);

await reverter.snapshot();
Expand Down Expand Up @@ -1312,7 +1311,7 @@ describe('Distribution', () => {
expect(await depositToken.balanceOf(distribution)).to.eq(wei(20));

await setNextTime(oneDay + oneDay);
await depositToken.setBalanceMultiplicator(wei(0.8, 25));
await depositToken.setTotalPooledEther(((await depositToken.totalPooledEther()) * 8n) / 10n);
expect(await depositToken.balanceOf(distribution)).to.eq(wei(16));

let tx = await distribution.withdraw(poolId, wei(999));
Expand All @@ -1330,7 +1329,7 @@ describe('Distribution', () => {
it('should revert if trying to withdraw zero', async () => {
await distribution.stake(poolId, wei(10));

await depositToken.setBalanceMultiplicator(wei(0.0001, 25));
await depositToken.setTotalPooledEther(wei(0.0001, 25));

await distribution.withdraw(poolId, wei(10));

Expand Down Expand Up @@ -1582,15 +1581,15 @@ describe('Distribution', () => {

await setTime(oneDay * 9999);

await depositToken.setBalanceMultiplicator(wei(0.5));
await depositToken.setTotalPooledEther(wei(0.5));

const overplus = await distribution.overplus();
expect(overplus).to.eq(0);
});
it('should return overplus if deposited token increased', async () => {
await distribution.stake(0, wei(1));

await depositToken.setBalanceMultiplicator(wei(2, 25));
await depositToken.setTotalPooledEther((await depositToken.totalPooledEther()) * 2n);

let overplus = await distribution.overplus();
expect(overplus).to.eq(wei(1));
Expand All @@ -1600,12 +1599,12 @@ describe('Distribution', () => {
overplus = await distribution.overplus();
expect(overplus).to.eq(wei(1));

await depositToken.setBalanceMultiplicator(wei(1, 25));
await depositToken.setTotalPooledEther((await depositToken.totalPooledEther()) / 2n);

overplus = await distribution.overplus();
expect(overplus).to.eq(0);

await depositToken.setBalanceMultiplicator(wei(5, 25));
await depositToken.setTotalPooledEther((await depositToken.totalPooledEther()) * 5n);

overplus = await distribution.overplus();
expect(overplus).to.eq(wei(5.5));
Expand All @@ -1626,7 +1625,7 @@ describe('Distribution', () => {

await distribution.stake(1, wei(1));

await depositToken.setBalanceMultiplicator(wei(2, 25));
await depositToken.setTotalPooledEther((await depositToken.totalPooledEther()) * 2n);

const overplus = await distribution.overplus();
expect(overplus).to.eq(wei(1));
Expand Down

0 comments on commit 07d8c1c

Please sign in to comment.