Merge pull request #5085 from BitGo/SC-498

Author: noel-bitgo

modules/bitgo/test/v2/fixtures/staking/stakingWallet.ts

@@ -20,7 +20,11 @@ export default {
       transactions: transactions,
     };
   },
-  transaction: function (status: string, buildParams?: PrebuildTransactionOptions): StakingTransaction {
+  transaction: function (
+    status: string,
+    buildParams?: PrebuildTransactionOptions,
+    setTxRequestId = true
+  ): StakingTransaction {
     const transaction: StakingTransaction = {
       id: '00566722-daef-40eb-b0ac-fa5402bbfe72',
       stakingRequestId: '8638284a-dab2-46b9-b07f-21109a6e7220',
@@ -32,11 +36,13 @@ export default {
       amount: '1234',
       pendingApprovalId: 'd99e3ae1-d2a6-4f57-87b6-d04c24854739',
       transferId: 'e4b482b0-54d5-474b-bb2b-c56ce8516b5e',
-      txRequestId: this.txRequestId,
     };
     if (buildParams) {
       transaction.buildParams = buildParams;
     }
+    if (setTxRequestId) {
+      transaction.txRequestId = this.txRequestId;
+    }
     return transaction;
   },
   buildParams: {

modules/bitgo/test/v2/unit/staking/stakingWalletCommon.ts

@@ -2,7 +2,7 @@ import * as nock from 'nock';
 import * as should from 'should';
 import fixtures from '../../fixtures/staking/stakingWallet';
 
-import { Enterprise, Environments, StakingRequest, StakingWallet, Wallet } from '@bitgo/sdk-core';
+import { Enterprise, Environments, StakingRequest, StakingWallet, TssUtils, Wallet } from '@bitgo/sdk-core';
 import { TestBitGo } from '@bitgo/sdk-test';
 import { BitGo } from '../../../../src';
 import * as sinon from 'sinon';
@@ -309,32 +309,73 @@ describe('Staking Wallet Common', function () {
 
   describe('prebuildSelfManagedStakingTransaction', function () {
     it('should prebuild self-managed staking transaction', async function () {
-      const transaction = fixtures.transaction('READY', fixtures.buildParams);
+      const transaction = fixtures.transaction('READY', fixtures.buildParams, false);
       nock(microservicesUri)
         .get(
           `/api/staking/v1/${stakingWallet.coin}/wallets/${stakingWallet.walletId}/requests/${transaction.stakingRequestId}/transactions/${transaction.id}`
         )
         .query({ expandBuildParams: true })
         .reply(200, transaction);
 
+      const deleteSignatureShares = sandbox.stub(TssUtils.prototype, 'deleteSignatureShares');
       const prebuildTransaction = sandbox.stub(Wallet.prototype, 'prebuildTransaction');
+      const build = sandbox.stub(StakingWallet.prototype, 'build');
       const txPrebuild = {
         walletId: stakingWallet.walletId,
         txHex: 'hex',
         buildParams: transaction.buildParams,
       };
       prebuildTransaction.resolves(txPrebuild);
       prebuildTransaction.calledOnceWithExactly(transaction.buildParams);
-
       const formattedParams = {
         ...fixtures.buildParams,
         coin: stakingWallet.coin,
         walletId: stakingWallet.walletId,
         walletType: stakingWallet.wallet.type(),
         preview: true,
       };
+      const stakingTransaction = await stakingWallet.prebuildSelfManagedStakingTransaction(transaction);
+      sandbox.assert.calledOnce(prebuildTransaction);
+      sandbox.assert.notCalled(build);
+      sandbox.assert.notCalled(deleteSignatureShares);
+
       const expected = await stakingWallet.wallet.prebuildTransaction(formattedParams);
+      stakingTransaction.should.deepEqual(expected);
+      should.exist(stakingTransaction);
+    });
+    it('should prebuild self-managed staking transaction - no build params', async function () {
+      const transaction = fixtures.transaction('READY', undefined, true);
+      nock(microservicesUri)
+        .get(
+          `/api/staking/v1/${stakingWallet.coin}/wallets/${stakingWallet.walletId}/requests/${transaction.stakingRequestId}/transactions/${transaction.id}`
+        )
+        .query({ expandBuildParams: true })
+        .reply(200, transaction);
+
+      const deleteSignatureShares = sandbox.stub(TssUtils.prototype, 'deleteSignatureShares');
+      const prebuildTransaction = sandbox.stub(Wallet.prototype, 'prebuildTransaction');
+      const build = sandbox.stub(StakingWallet.prototype, 'build');
+      const txPrebuild = {
+        walletId: stakingWallet.walletId,
+        txHex: 'hex',
+      };
+      prebuildTransaction.resolves(txPrebuild);
+      if (transaction.txRequestId) {
+        deleteSignatureShares.calledOnceWithExactly(transaction.txRequestId);
+      }
+      const formattedParams = {
+        ...fixtures.buildParams,
+        coin: stakingWallet.coin,
+        walletId: stakingWallet.walletId,
+        walletType: stakingWallet.wallet.type(),
+        preview: true,
+      };
       const stakingTransaction = await stakingWallet.prebuildSelfManagedStakingTransaction(transaction);
+      sandbox.assert.calledOnce(prebuildTransaction);
+      sandbox.assert.notCalled(build);
+      sandbox.assert.calledOnce(deleteSignatureShares);
+
+      const expected = await stakingWallet.wallet.prebuildTransaction(formattedParams);
 
       stakingTransaction.should.deepEqual(expected);
       should.exist(stakingTransaction);

modules/sdk-core/src/bitgo/staking/stakingWallet.ts

@@ -218,19 +218,25 @@ export class StakingWallet implements IStakingWallet {
 
   /**
    * Create prebuilt staking transaction.
+   *
+   * for transactions with tx request id (TSS transactions), we need to delete signature shares before creating prebuild transaction
+   * we only need to get transaction build params if they exist to pre build
+   *
    * @param transaction
    */
   async prebuildSelfManagedStakingTransaction(transaction: StakingTransaction): Promise<PrebuildTransactionResult> {
-    const builtStakingTransaction = await this.build(transaction);
-    const buildParams = builtStakingTransaction.transaction.buildParams;
+    if (transaction.txRequestId) {
+      await this.tssUtil.deleteSignatureShares(transaction.txRequestId);
+    }
+    const buildParams = (await this.expandBuildParams(transaction)).buildParams;
     const formattedParams = {
       ...buildParams,
       coin: this.coin,
       walletId: this.walletId,
       walletType: this.wallet.type(),
       preview: true,
     };
-    return await this.wallet.prebuildTransaction(formattedParams);
+    return await (await this.getWalletForBuildingAndSigning()).prebuildTransaction(formattedParams);
   }
 
   /**