add unit tests

contracts/Controller.sol

@@ -106,8 +106,8 @@ contract Controller is ReentrancyGuard, Ownable {
         address indexed otoken,
         address indexed from,
         address indexed exerciser,
+        address collateralAsset,
         uint256 otokenBurned,
-        uint256 cashValue,
         uint256 payout
     );
 
@@ -540,21 +540,31 @@ contract Controller is ReentrancyGuard, Ownable {
 
         require(isPriceFinalized(_args.otoken), "Controller: otoken underlying asset price is not finalized yet");
 
-        address calculatorModule = AddressBookInterface(addressBook).getMarginCalculator();
-        MarginCalculatorInterface calculator = MarginCalculatorInterface(calculatorModule);
-
-        uint256 cashValue = calculator.getExpiredCashValue(_args.otoken);
-
-        uint256 payout = cashValue.mul(_args.amount);
+        uint256 payout = getPayout(_args.otoken, _args.amount);
 
         otoken.burnOtoken(msg.sender, _args.amount);
 
         address marginPoolModule = AddressBookInterface(addressBook).getMarginPool();
         MarginPoolInterface marginPool = MarginPoolInterface(marginPoolModule);
 
-        marginPool.transferToUser(_args.otoken, _args.exerciser, payout);
+        marginPool.transferToUser(otoken.collateralAsset(), _args.exerciser, payout);
+
+        emit Exercise(_args.otoken, msg.sender, _args.exerciser, otoken.collateralAsset(), _args.amount, payout);
+    }
+
+    /**
+     * @notice get Otoken payout after expiry
+     * @param _otoken Otoken address
+     * @param _amount amount of Otoken
+     * @return payout = cashValue * amount
+     */
+    function getPayout(address _otoken, uint256 _amount) internal view returns (uint256) {
+        address calculatorModule = AddressBookInterface(addressBook).getMarginCalculator();
+        MarginCalculatorInterface calculator = MarginCalculatorInterface(calculatorModule);
+
+        uint256 cashValue = calculator.getExpiredCashValue(_otoken);
 
-        emit Exercise(_args.otoken, msg.sender, _args.exerciser, _args.amount, cashValue, payout);
+        return cashValue.mul(_amount).div(1e18);
     }
 
     /**

contracts/MarginCalculator.sol

@@ -296,7 +296,6 @@ contract MarginCalculator is Initializable {
     function _getUnderlyingPrice(address _otoken) internal view returns (uint256 price, bool isFinalized) {
         OracleInterface oracle = OracleInterface(AddressBookInterface(addressBook).getOracle());
         OtokenInterface otoken = OtokenInterface(_otoken);
-        // return (otoken.expiryTimestamp(), true);
         return oracle.getExpiryPrice(otoken.underlyingAsset(), otoken.expiryTimestamp());
     }
 

contracts/mocks/MockMarginCalculator.sol

@@ -6,6 +6,8 @@ pragma solidity =0.6.10;
 pragma experimental ABIEncoderV2;
 
 import {OtokenInterface} from "../interfaces/OtokenInterface.sol";
+import {OracleInterface} from "../interfaces/OracleInterface.sol";
+import {AddressBookInterface} from "../interfaces/AddressBookInterface.sol";
 import {MarginAccount} from "../libs/MarginAccount.sol";
 import {FixedPointInt256} from "../libs/FixedPointInt256.sol";
 import {SignedConverter} from "../libs/SignedConverter.sol";
@@ -15,8 +17,11 @@ contract MockMarginCalculator {
     using SafeMath for uint256;
     using FixedPointInt256 for FixedPointInt256.FixedPointInt;
 
-    mapping(address => uint256) public price;
-    mapping(address => bool) public isPriceFinalized;
+    address addressBook;
+
+    constructor(address _addressBook) public {
+        addressBook = _addressBook;
+    }
 
     function getExpiredCashValue(address _otoken) public view returns (uint256) {
         require(_otoken != address(0), "MarginCalculator: Invalid token address.");
@@ -206,7 +211,10 @@ contract MockMarginCalculator {
     }
 
     function _getUnderlyingPrice(address _otoken) internal view returns (uint256, bool) {
-        return (price[_otoken], isPriceFinalized[_otoken]);
+        OracleInterface oracle = OracleInterface(AddressBookInterface(addressBook).getOracle());
+        OtokenInterface otoken = OtokenInterface(_otoken);
+
+        return oracle.getExpiryPrice(otoken.underlyingAsset(), otoken.expiryTimestamp());
     }
 
     function _uint256ToFixedPointInt(uint256 _num) internal pure returns (FixedPointInt256.FixedPointInt memory) {

test/controller.test.ts

@@ -37,12 +37,10 @@ enum ActionType {
   Call,
 }
 
-contract('Controller', ([owner, accountOwner1, accountOperator1, random]) => {
+contract('Controller', ([owner, accountOwner1, accountOperator1, holder1, random]) => {
   // ERC20 mock
   let usdc: MockERC20Instance
   let weth: MockERC20Instance
-  // Otoken mock
-  let otoken: MockOtokenInstance
   // Oracle module
   let oracle: MockOracleInstance
   // calculator module
@@ -60,25 +58,12 @@ contract('Controller', ([owner, accountOwner1, accountOperator1, random]) => {
     // ERC20 deployment
     usdc = await MockERC20.new('USDC', 'USDC', 8)
     weth = await MockERC20.new('WETH', 'WETH', 18)
-    // Otoken deployment
-    otoken = await MockOtoken.new()
     // addressbook deployment
     addressBook = await MockAddressBook.new()
-    // init otoken
-    await otoken.init(
-      addressBook.address,
-      weth.address,
-      usdc.address,
-      usdc.address,
-      new BigNumber(200).times(new BigNumber(10).exponentiatedBy(18)),
-      1753776000, // 07/29/2025 @ 8:00am (UTC)
-      true,
-    )
-
     // deploy Oracle module
     oracle = await MockOracle.new(addressBook.address, {from: owner})
     // calculator deployment
-    calculator = await MockMarginCalculator.new()
+    calculator = await MockMarginCalculator.new(addressBook.address)
     // margin pool deployment
     marginPool = await MarginPool.new(addressBook.address)
     // whitelist module
@@ -2158,8 +2143,6 @@ contract('Controller', ([owner, accountOwner1, accountOperator1, random]) => {
         const vaultBefore = await controller.getVault(accountOwner1, vaultCounter)
 
         const shortOtokenToBurn = new BigNumber(vaultBefore.shortAmounts[0])
-        console.log(shortOtokenToBurn.toString())
-        console.log(new BigNumber(await shortOtoken.balanceOf(accountOwner1)).toString())
         const actionArgs = [
           {
             actionType: ActionType.BurnShortOption,
@@ -2301,6 +2284,221 @@ contract('Controller', ([owner, accountOwner1, accountOperator1, random]) => {
     })
   })
 
+  describe('Exercise', () => {
+    let shortOtoken: MockOtokenInstance
+
+    before(async () => {
+      const expiryTime = new BigNumber(60 * 60 * 24) // after 1 day
+
+      shortOtoken = await MockOtoken.new()
+      // init otoken
+      await shortOtoken.init(
+        addressBook.address,
+        weth.address,
+        usdc.address,
+        usdc.address,
+        new BigNumber(200).times(new BigNumber(10).exponentiatedBy(18)),
+        new BigNumber(await time.latest()).plus(expiryTime),
+        true,
+      )
+      // whitelist short otoken to be used in the protocol
+      await whitelist.whitelistOtoken(shortOtoken.address, {from: owner})
+      // give free money
+      await usdc.mint(accountOwner1, new BigNumber('1000000'))
+      await usdc.mint(accountOperator1, new BigNumber('1000000'))
+      await usdc.mint(random, new BigNumber('1000000'))
+      // open new vault, mintnaked short, sell it to holder 1
+      const vaultCounter = new BigNumber(await controller.getAccountVaultCounter(accountOwner1)).plus(1)
+      const collateralToDeposit = new BigNumber(await shortOtoken.strikePrice()).dividedBy(1e18)
+      const amountToMint = new BigNumber('1')
+      const actionArgs = [
+        {
+          actionType: ActionType.OpenVault,
+          owner: accountOwner1,
+          sender: accountOwner1,
+          asset: ZERO_ADDR,
+          vaultId: vaultCounter.toNumber(),
+          amount: '0',
+          index: '0',
+          data: ZERO_ADDR,
+        },
+        {
+          actionType: ActionType.MintShortOption,
+          owner: accountOwner1,
+          sender: accountOwner1,
+          asset: shortOtoken.address,
+          vaultId: vaultCounter.toNumber(),
+          amount: amountToMint.toNumber(),
+          index: '0',
+          data: ZERO_ADDR,
+        },
+        {
+          actionType: ActionType.DepositCollateral,
+          owner: accountOwner1,
+          sender: accountOwner1,
+          asset: usdc.address,
+          vaultId: vaultCounter.toNumber(),
+          amount: collateralToDeposit.toNumber(),
+          index: '0',
+          data: ZERO_ADDR,
+        },
+      ]
+      const marginPoolBalanceBefore = new BigNumber(await usdc.balanceOf(marginPool.address))
+      const senderBalanceBefore = new BigNumber(await usdc.balanceOf(accountOwner1))
+      const senderShortBalanceBefore = new BigNumber(await shortOtoken.balanceOf(accountOwner1))
+
+      await usdc.approve(marginPool.address, collateralToDeposit, {from: accountOwner1})
+      await controller.operate(actionArgs, {from: accountOwner1})
+
+      const marginPoolBalanceAfter = new BigNumber(await usdc.balanceOf(marginPool.address))
+      const senderBalanceAfter = new BigNumber(await usdc.balanceOf(accountOwner1))
+      const senderShortBalanceAfter = new BigNumber(await shortOtoken.balanceOf(accountOwner1))
+      const vaultAfter = await controller.getVault(accountOwner1, vaultCounter)
+      assert.equal(
+        marginPoolBalanceAfter.minus(marginPoolBalanceBefore).toString(),
+        collateralToDeposit.toString(),
+        'Margin pool collateral asset balance mismatch',
+      )
+      assert.equal(
+        senderBalanceBefore.minus(senderBalanceAfter).toString(),
+        collateralToDeposit.toString(),
+        'Sender collateral asset balance mismatch',
+      )
+      assert.equal(vaultAfter.collateralAssets.length, 1, 'Vault collateral asset array length mismatch')
+      assert.equal(vaultAfter.shortOtokens.length, 1, 'Vault short otoken array length mismatch')
+      assert.equal(
+        vaultAfter.collateralAssets[0],
+        usdc.address,
+        'Collateral asset address deposited into vault mismatch',
+      )
+      assert.equal(
+        vaultAfter.shortOtokens[0],
+        shortOtoken.address,
+        'Short otoken address deposited into vault mismatch',
+      )
+      assert.equal(
+        senderShortBalanceAfter.minus(senderShortBalanceBefore).toString(),
+        amountToMint.toString(),
+        'Short otoken amount minted mismatch',
+      )
+
+      // transger minted short ototken to hodler`
+      shortOtoken.transfer(holder1, amountToMint, {from: accountOwner1})
+    })
+
+    it('should revert exercising un-expired otoken', async () => {
+      const shortAmountToBurn = new BigNumber('1')
+      const actionArgs = [
+        {
+          actionType: ActionType.Exercise,
+          owner: ZERO_ADDR,
+          sender: holder1,
+          asset: shortOtoken.address,
+          vaultId: '0',
+          amount: shortAmountToBurn.toNumber(),
+          index: '0',
+          data: ZERO_ADDR,
+        },
+      ]
+
+      assert.equal(await controller.isExpired(shortOtoken.address), false, 'Short otoken is already expired')
+
+      await expectRevert(
+        controller.operate(actionArgs, {from: holder1}),
+        'Controller: can not exercise un-expired otoken',
+      )
+    })
+
+    it('should revert exercising after expiry, when price is not finalized yet', async () => {
+      // past time after expiry
+      await time.increase(60 * 61 * 24) // increase time with one hour in seconds
+      // set price in Oracle Mock, 150$ at expiry, expire ITM
+      await oracle.setExpiryPrice(
+        await shortOtoken.underlyingAsset(),
+        new BigNumber(await shortOtoken.expiryTimestamp()),
+        new BigNumber(150).times(new BigNumber(10).exponentiatedBy(18)),
+      )
+      // set it as not finalized in mock
+      await oracle.setIsFinalized(
+        await shortOtoken.underlyingAsset(),
+        new BigNumber(await shortOtoken.expiryTimestamp()),
+        false,
+      )
+
+      const shortAmountToBurn = new BigNumber('1')
+      const actionArgs = [
+        {
+          actionType: ActionType.Exercise,
+          owner: ZERO_ADDR,
+          sender: holder1,
+          asset: shortOtoken.address,
+          vaultId: '0',
+          amount: shortAmountToBurn.toNumber(),
+          index: '0',
+          data: ZERO_ADDR,
+        },
+      ]
+
+      assert.equal(await controller.isExpired(shortOtoken.address), true, 'Short otoken is not expired yet')
+
+      await expectRevert(
+        controller.operate(actionArgs, {from: holder1}),
+        'Controller: otoken underlying asset price is not finalized yet',
+      )
+    })
+
+    it('should exercise after expiry + price is finalized', async () => {
+      // set it as finalized in mock
+      await oracle.setIsFinalized(
+        await shortOtoken.underlyingAsset(),
+        new BigNumber(await shortOtoken.expiryTimestamp()),
+        true,
+      )
+
+      const shortAmountToBurn = new BigNumber('1')
+      const actionArgs = [
+        {
+          actionType: ActionType.Exercise,
+          owner: ZERO_ADDR,
+          sender: holder1,
+          asset: shortOtoken.address,
+          vaultId: '0',
+          amount: shortAmountToBurn.toNumber(),
+          index: '0',
+          data: ZERO_ADDR,
+        },
+      ]
+      assert.equal(await controller.isExpired(shortOtoken.address), true, 'Short otoken is not expired yet')
+
+      const payout = new BigNumber('50')
+      const marginPoolBalanceBefore = new BigNumber(await usdc.balanceOf(marginPool.address))
+      const senderBalanceBefore = new BigNumber(await usdc.balanceOf(holder1))
+      const senderShortBalanceBefore = new BigNumber(await shortOtoken.balanceOf(holder1))
+
+      controller.operate(actionArgs, {from: holder1})
+
+      const marginPoolBalanceAfter = new BigNumber(await usdc.balanceOf(marginPool.address))
+      const senderBalanceAfter = new BigNumber(await usdc.balanceOf(holder1))
+      const senderShortBalanceAfter = new BigNumber(await shortOtoken.balanceOf(holder1))
+
+      assert.equal(
+        marginPoolBalanceBefore.minus(marginPoolBalanceAfter).toString(),
+        payout.toString(),
+        'Margin pool collateral asset balance mismatch',
+      )
+      assert.equal(
+        senderBalanceAfter.minus(senderBalanceBefore).toString(),
+        payout.toString(),
+        'Sender collateral asset balance mismatch',
+      )
+      assert.equal(
+        senderShortBalanceBefore.minus(senderShortBalanceAfter).toString(),
+        shortAmountToBurn.toString(),
+        ' Burned short otoken amount mismatch',
+      )
+    })
+  })
+
   describe('Check if price is finalized', () => {
     let expiredOtoken: MockOtokenInstance
     let expiry: BigNumber
@@ -2369,6 +2567,17 @@ contract('Controller', ([owner, accountOwner1, accountOperator1, random]) => {
 
   describe('Expiry', () => {
     it('should return false for non expired otoken', async () => {
+      const otoken: MockOtokenInstance = await MockOtoken.new()
+      await otoken.init(
+        addressBook.address,
+        weth.address,
+        usdc.address,
+        usdc.address,
+        new BigNumber(200).times(new BigNumber(10).exponentiatedBy(18)),
+        new BigNumber(await time.latest()).plus(60000 * 60000),
+        true,
+      )
+
       assert.equal(await controller.isExpired(otoken.address), false, 'Otoken expiry check mismatch')
     })