From ffe28b17440349f9de46aac9853f30d0a61bd060 Mon Sep 17 00:00:00 2001
From: Marc-Aurele Besner
 <82244926+marc-aurele-besner@users.noreply.github.com>
Date: Wed, 12 Jun 2024 16:49:54 -0400
Subject: [PATCH] add/improve staking tests

---
 .../auto-consensus/__test__/staking.test.ts   | 82 ++++++++++++++++++-
 1 file changed, 78 insertions(+), 4 deletions(-)

diff --git a/packages/auto-consensus/__test__/staking.test.ts b/packages/auto-consensus/__test__/staking.test.ts
index 9bb24902..1a4fab58 100644
--- a/packages/auto-consensus/__test__/staking.test.ts
+++ b/packages/auto-consensus/__test__/staking.test.ts
@@ -2,9 +2,16 @@ import { ActivateWalletInput, activateWallet } from '@autonomys/auto-utils'
 import { mnemonicGenerate } from '@polkadot/util-crypto'
 import { address } from '../src/address'
 import { balance } from '../src/balances'
-import { deregisterOperator, nominateOperator, registerOperator } from '../src/staking'
+import {
+  deregisterOperator,
+  nominateOperator,
+  registerOperator,
+  unlockFunds,
+  withdrawStake,
+} from '../src/staking'
 import { transfer } from '../src/transfer'
 import {
+  events,
   setup,
   signAndSendTx,
   sudo,
@@ -47,7 +54,7 @@ describe('Verify staking functions', () => {
           minimumNominatorStake,
           nominationTax,
         }
-        await signAndSendTx(sender, await registerOperator(txInput))
+        await signAndSendTx(sender, await registerOperator(txInput), [events.operatorRegistered])
         const findOperator = await verifyOperatorRegistration(txInput)
 
         const _balanceSenderEnd = await balance(api, address(sender.address))
@@ -123,7 +130,7 @@ describe('Verify staking functions', () => {
           minimumNominatorStake,
           nominationTax,
         }
-        await signAndSendTx(sender, await registerOperator(txInput))
+        await signAndSendTx(sender, await registerOperator(txInput), [events.operatorRegistered])
         await verifyOperatorRegistration(txInput)
 
         const _balanceSenderEnd = await balance(api, address(sender.address))
@@ -136,6 +143,72 @@ describe('Verify staking functions', () => {
       }, 30000)
     })
 
+    describe('Test withdrawStake()', () => {
+      test('Check Operator can withdrawStake after registration', async () => {
+        const mnemonicOperator = mnemonicGenerate()
+        const { api, accounts } = await activateWallet({
+          ...TEST_NETWORK,
+          uri: ALICE_URI,
+        } as ActivateWalletInput)
+        const { accounts: operatorAccounts } = await activateWallet({
+          ...TEST_NETWORK,
+          uri: mnemonicOperator,
+        } as ActivateWalletInput)
+        expect(accounts.length).toBeGreaterThan(0)
+        expect(accounts[0].address).toEqual(ALICE_ADDRESS)
+
+        const sender = accounts[0]
+
+        const _balanceSenderStart = await balance(api, address(sender.address))
+        expect(_balanceSenderStart.free).toBeGreaterThan(BigInt(0))
+
+        // Transfer some funds to the operator
+        const amountToTransfer = '10000000000000000000'
+        await signAndSendTx(
+          sender,
+          await transfer(api, operatorAccounts[0].address, amountToTransfer),
+        )
+
+        const domainId = '0'
+        const amountToStake = '100000000000000000000'
+        const minimumNominatorStake = '1000000000000000000'
+        const nominationTax = '5'
+        const txInput = {
+          api,
+          senderAddress: ALICE_ADDRESS,
+          Operator: operatorAccounts[0],
+          domainId,
+          amountToStake,
+          minimumNominatorStake,
+          nominationTax,
+        }
+        await signAndSendTx(sender, await registerOperator(txInput), [events.operatorRegistered])
+        await verifyOperatorRegistration(txInput)
+
+        const _balanceSenderEnd = await balance(api, address(sender.address))
+        expect(_balanceSenderEnd.free).toBeLessThan(
+          _balanceSenderStart.free - BigInt(amountToStake),
+        )
+
+        await sudo(api, sender, await api.tx.domains.forceStakingEpochTransition(domainId))
+        const findOperator = await verifyOperatorRegistrationFinal(txInput)
+
+        await sudo(api, sender, await api.tx.domains.forceStakingEpochTransition(domainId))
+
+        if (findOperator) {
+          await signAndSendTx(
+            operatorAccounts[0],
+            await withdrawStake({
+              api,
+              operatorId: findOperator.operatorId,
+              shares: findOperator.operatorDetails.currentTotalShares,
+            }),
+            [events.withdrawStake],
+          )
+        }
+      }, 60000)
+    })
+
     describe('Test deregisterOperator()', () => {
       test('Check Operator can deregisterOperator after registration', async () => {
         const mnemonicOperator = mnemonicGenerate()
@@ -175,7 +248,7 @@ describe('Verify staking functions', () => {
           minimumNominatorStake,
           nominationTax,
         }
-        await signAndSendTx(sender, await registerOperator(txInput))
+        await signAndSendTx(sender, await registerOperator(txInput), [events.operatorRegistered])
         await verifyOperatorRegistration(txInput)
 
         const _balanceSenderEnd = await balance(api, address(sender.address))
@@ -193,6 +266,7 @@ describe('Verify staking functions', () => {
               api,
               operatorId: findOperator.operatorId,
             }),
+            [events.operatorDeRegistered],
           )
         }
       }, 30000)