Merge pull request #103 from invariant-labs/sdk-protocol-fee-tests

add withdraw protocol fee test

sdk/src/erc20.ts

@@ -140,9 +140,12 @@ export class Erc20 {
   async balanceOf(addressHash: Key, address: string): Promise<bigint> {
     const balanceKey = new Uint8Array([...BALANCES, addressHash, ...hexToBytes(address)])
 
-    const response = await this.contract.queryContractDictionary('state', hash(balanceKey))
-
-    return BigInt(response.data._hex)
+    try {
+      const response = await this.contract.queryContractDictionary('state', hash(balanceKey))
+      return BigInt(response.data._hex)
+    } catch (e) {
+      return 0n
+    }
   }
 
   async allowance(

sdk/src/index.ts

@@ -2,7 +2,13 @@ import { ALICE, BOB, LOCAL_NODE_URL, TEST, TESTNET_NODE_URL } from './consts'
 import { Key, Network } from './enums'
 import { Erc20 } from './erc20'
 import { Invariant } from './invariant'
-import { callWasm, createAccountKeys, initCasperClient, loadWasm } from './utils'
+import {
+  callWasm,
+  createAccountKeys,
+  getAccountHashFromKey,
+  initCasperClient,
+  loadWasm
+} from './utils'
 const main = async () => {
   const createKeys = false
   const wasm = await loadWasm()
@@ -16,17 +22,17 @@ const main = async () => {
   const isLocal = true
 
   let account = ALICE
-  let accountAddress = account.publicKey.toAccountHashStr().replace('account-hash-', '')
+  let accountAddress = getAccountHashFromKey(account)
   const dummy = BOB
   /* eslint-disable-next-line @typescript-eslint/no-unused-vars */
-  const dummyAddress = dummy.publicKey.toAccountHashStr().replace('account-hash-', '')
+  const dummyAddress = getAccountHashFromKey(dummy)
   let network = Network.Local
   let nodeUrl = LOCAL_NODE_URL
 
   if (!isLocal) {
     account = TEST
     /* eslint-disable-next-line @typescript-eslint/no-unused-vars */
-    accountAddress = account.publicKey.toAccountHashStr().replace('account-hash-', '')
+    accountAddress = getAccountHashFromKey(account)
     network = Network.Testnet
     nodeUrl = TESTNET_NODE_URL
   }

sdk/src/invariant.ts

@@ -31,7 +31,7 @@ import {
   decodePositionLength,
   decodeTick
 } from './decoder'
-import { Network } from './enums'
+import { Key, Network } from './enums'
 import { bigintToByteArray, encodePoolKey, hash } from './parser'
 import {
   callWasm,
@@ -202,10 +202,19 @@ export class Invariant {
     )
   }
 
-  async changeFeeReceiver(signer: Keys.AsymmetricKey, poolKey: PoolKey, newFeeReceiver: string) {
+  async changeFeeReceiver(
+    signer: Keys.AsymmetricKey,
+    poolKey: PoolKey,
+    newFeeReceiverHash: Key,
+    newFeeReceiver: string
+  ) {
     const token0Key = new CLByteArray(decodeBase16(poolKey.tokenX))
     const token1Key = new CLByteArray(decodeBase16(poolKey.tokenY))
-    const feeReceiverKey = new CLByteArray(decodeBase16(newFeeReceiver))
+    const newFeeReceiverBytes = new Uint8Array([
+      newFeeReceiverHash,
+      ...decodeBase16(newFeeReceiver)
+    ])
+    const newFeeReceiverKey = new CLByteArray(newFeeReceiverBytes)
 
     return await sendTx(
       this.contract,
@@ -219,7 +228,7 @@ export class Invariant {
         token_1: CLValueBuilder.key(token1Key),
         fee: CLValueBuilder.u128(BigNumber.from(poolKey.feeTier.fee.v)),
         tick_spacing: CLValueBuilder.u32(integerSafeCast(poolKey.feeTier.tickSpacing)),
-        fee_receiver: CLValueBuilder.key(feeReceiverKey)
+        fee_receiver: newFeeReceiverKey
       }
     )
   }

sdk/src/utils.ts

@@ -294,3 +294,7 @@ export const extractContractHash = (contractPackageHash: string): string => {
 export const extractContractPackageHash = (contractPackage: ContractPackageJson): string => {
   return contractPackage.versions[0].contractHash.replace('contract-', '')
 }
+
+export const getAccountHashFromKey = (key: Keys.AsymmetricKey): string => {
+  return key.publicKey.toAccountHashStr().replace('account-hash-', '')
+}

sdk/tests/erc20.test.ts

@@ -2,12 +2,12 @@ import { ALICE, BOB, LOCAL_NODE_URL } from '../src/consts'
 import { Key, Network } from '../src/enums'
 import { Erc20 } from '../src/erc20'
 import { loadChai } from '../src/testUtils'
-import { initCasperClient } from '../src/utils'
+import { getAccountHashFromKey, initCasperClient } from '../src/utils'
 
 const client = initCasperClient(LOCAL_NODE_URL)
 let erc20: Erc20
-const aliceAddress = ALICE.publicKey.toAccountHashStr().replace('account-hash-', '')
-const bobAddress = BOB.publicKey.toAccountHashStr().replace('account-hash-', '')
+const aliceAddress = getAccountHashFromKey(ALICE)
+const bobAddress = getAccountHashFromKey(BOB)
 
 describe('erc20', () => {
   beforeEach(async () => {

sdk/tests/math.test.ts

@@ -9,7 +9,7 @@ import {
   loadChai,
   positionEquals
 } from '../src/testUtils'
-import { callWasm, initCasperClient, loadWasm } from '../src/utils'
+import { callWasm, getAccountHashFromKey, initCasperClient, loadWasm } from '../src/utils'
 
 let hashes: {
   invariant: { loadHash: string; packageHash: string }
@@ -21,7 +21,7 @@ describe('test get liquidity by x', () => {
   const client = initCasperClient(LOCAL_NODE_URL)
   const deployer = ALICE
   const positionOwner = BOB
-  const positionOwnerHash = positionOwner.publicKey.toAccountHashStr().replace('account-hash-', '')
+  const positionOwnerHash = getAccountHashFromKey(positionOwner)
   const network = Network.Local
   const providedAmount = { v: 430000n }
   let feeTier: FeeTier
@@ -168,7 +168,7 @@ describe('test get liquidity by y', () => {
   const client = initCasperClient(LOCAL_NODE_URL)
   const deployer = ALICE
   const positionOwner = BOB
-  const positionOwnerHash = positionOwner.publicKey.toAccountHashStr().replace('account-hash-', '')
+  const positionOwnerHash = getAccountHashFromKey(positionOwner)
   const network = Network.Local
   const providedAmount = { v: 47600000000n }
   let feeTier: FeeTier

sdk/tests/protocol-fee.test.ts

@@ -0,0 +1,171 @@
+import { FeeTier, PoolKey } from 'invariant-cspr-wasm'
+import { ALICE, BOB, LOCAL_NODE_URL } from '../src/consts'
+import { Key, Network } from '../src/enums'
+import { Erc20 } from '../src/erc20'
+import { Invariant } from '../src/invariant'
+import { loadChai } from '../src/testUtils'
+import { callWasm, getAccountHashFromKey, initCasperClient, loadWasm } from '../src/utils'
+
+let wasm: typeof import('invariant-cspr-wasm')
+let chai: typeof import('chai')
+
+const client = initCasperClient(LOCAL_NODE_URL)
+let erc20: Erc20
+let invariant: Invariant
+let invariantContractPackage: string
+let token0Address: string
+let token1Address: string
+let token0ContractPackage: string
+let token1ContractPackage: string
+const aliceAddress = getAccountHashFromKey(ALICE)
+const bobAddress = getAccountHashFromKey(BOB)
+
+const fee = 1000000000n
+const tickSpacing = 1n
+
+let feeTier: FeeTier
+let poolKey: PoolKey
+
+describe('protocol fee', () => {
+  before(async () => {
+    wasm = await loadWasm()
+    chai = await loadChai()
+  })
+
+  beforeEach(async () => {
+    const [token0ContractPackageHash, token0ContractHash] = await Erc20.deploy(
+      client,
+      Network.Local,
+      ALICE,
+      'erc20-1',
+      1000000000n,
+      'Coin',
+      'COIN',
+      12n,
+      150000000000n
+    )
+    const [token1ContractPackageHash, token1ContractHash] = await Erc20.deploy(
+      client,
+      Network.Local,
+      ALICE,
+      'erc20-2',
+      1000000000n,
+      'Coin',
+      'COIN',
+      12n,
+      150000000000n
+    )
+    const [invariantContractPackageHash, invariantContractHash] = await Invariant.deploy(
+      client,
+      Network.Local,
+      ALICE,
+      fee,
+      600000000000n
+    )
+
+    token0Address = token0ContractHash
+    token1Address = token1ContractHash
+    token0ContractPackage = token0ContractPackageHash
+    token1ContractPackage = token1ContractPackageHash
+    invariantContractPackage = invariantContractPackageHash
+
+    erc20 = await Erc20.load(client, Network.Local, token0ContractHash)
+    await erc20.approve(ALICE, Key.Hash, invariantContractPackage, fee)
+    erc20 = await Erc20.load(client, Network.Local, token1ContractHash)
+    await erc20.approve(ALICE, Key.Hash, invariantContractPackage, fee)
+
+    invariant = await Invariant.load(client, invariantContractHash, Network.Local)
+
+    feeTier = await callWasm(wasm.newFeeTier, { v: fee }, tickSpacing)
+    poolKey = await callWasm(wasm.newPoolKey, token0ContractPackage, token1ContractPackage, feeTier)
+
+    await invariant.addFeeTier(ALICE, feeTier)
+
+    await invariant.createPool(ALICE, poolKey, { v: 1000000000000000000000000n })
+
+    await invariant.createPosition(
+      ALICE,
+      poolKey,
+      -10n,
+      10n,
+      { v: 10000000000000n },
+      { v: 1000000000000000000000000n },
+      { v: 1000000000000000000000000n }
+    )
+
+    await invariant.swap(ALICE, poolKey, true, { v: 4999n }, true, { v: 999505344804856076727628n })
+  })
+
+  it('should withdraw protocol fee', async () => {
+    const feeTier = await callWasm(wasm.newFeeTier, { v: fee }, tickSpacing)
+    const poolKey = await callWasm(
+      wasm.newPoolKey,
+      token0ContractPackage,
+      token1ContractPackage,
+      feeTier
+    )
+
+    erc20.setContractHash(token0Address)
+    const token0Before = await erc20.balanceOf(Key.Account, aliceAddress)
+    erc20.setContractHash(token1Address)
+    const token1Before = await erc20.balanceOf(Key.Account, aliceAddress)
+
+    const poolBefore = await invariant.getPool(poolKey)
+    chai.assert.deepEqual(poolBefore.feeProtocolTokenX, { v: 1n })
+    chai.assert.deepEqual(poolBefore.feeProtocolTokenY, { v: 0n })
+
+    await invariant.withdrawProtocolFee(ALICE, poolKey)
+
+    const poolAfter = await invariant.getPool(poolKey)
+    chai.assert.deepEqual(poolAfter.feeProtocolTokenX, { v: 0n })
+    chai.assert.deepEqual(poolAfter.feeProtocolTokenY, { v: 0n })
+
+    erc20.setContractHash(token0Address)
+    const token0After = await erc20.balanceOf(Key.Account, aliceAddress)
+    erc20.setContractHash(token1Address)
+    const token1After = await erc20.balanceOf(Key.Account, aliceAddress)
+
+    if (await callWasm(wasm.isTokenX, token0ContractPackage, token1ContractPackage)) {
+      chai.assert.equal(token0After, token0Before + 1n)
+      chai.assert.equal(token1After, token1Before)
+    } else {
+      chai.assert.equal(token0After, token0Before)
+      chai.assert.equal(token1After, token1Before + 1n)
+    }
+  })
+
+  it('should change fee receiver', async () => {
+    await invariant.changeFeeReceiver(ALICE, poolKey, Key.Account, bobAddress)
+
+    erc20.setContractHash(token0Address)
+    const token0Before = await erc20.balanceOf(Key.Account, bobAddress)
+    erc20.setContractHash(token1Address)
+    const token1Before = await erc20.balanceOf(Key.Account, bobAddress)
+
+    const poolBefore = await invariant.getPool(poolKey)
+    chai.assert.deepEqual(poolBefore.feeProtocolTokenX, { v: 1n })
+    chai.assert.deepEqual(poolBefore.feeProtocolTokenY, { v: 0n })
+
+    const withdrawProtocolFeeResult = await invariant.withdrawProtocolFee(ALICE, poolKey)
+    chai.assert.notEqual(withdrawProtocolFeeResult.execution_results[0].result.Failure, undefined)
+
+    await invariant.withdrawProtocolFee(BOB, poolKey)
+
+    const poolAfter = await invariant.getPool(poolKey)
+    chai.assert.deepEqual(poolAfter.feeProtocolTokenX, { v: 0n })
+    chai.assert.deepEqual(poolAfter.feeProtocolTokenY, { v: 0n })
+
+    erc20.setContractHash(token0Address)
+    const token0After = await erc20.balanceOf(Key.Account, bobAddress)
+    erc20.setContractHash(token1Address)
+    const token1After = await erc20.balanceOf(Key.Account, bobAddress)
+
+    if (await callWasm(wasm.isTokenX, token0ContractPackage, token1ContractPackage)) {
+      chai.assert.equal(token0After, token0Before + 1n)
+      chai.assert.equal(token1After, token1Before)
+    } else {
+      chai.assert.equal(token0After, token0Before)
+      chai.assert.equal(token1After, token1Before + 1n)
+    }
+  })
+})