Simplify remove liquidity types

lib/modules/pool/actions/LiquidityActionHelpers.ts

@@ -10,6 +10,7 @@ import { toPoolStateInput } from '../pool.helpers'
 import { Pool } from '../usePool'
 import { HumanAmountIn } from './liquidity-types'
 import { GqlToken } from '@/lib/shared/services/api/generated/graphql'
+import { SentryError } from '@/lib/shared/utils/errors'
 
 // Null object used to avoid conditional checks during hook loading state
 const NullPool: Pool = {
@@ -106,3 +107,20 @@ export const hasValidHumanAmounts = (humanAmountsIn: HumanAmountIn[]) =>
 export function toHumanAmount(tokenAmount: TokenAmount): HumanAmount {
   return formatUnits(tokenAmount.amount, tokenAmount.token.decimals) as HumanAmount
 }
+
+export function ensureLastQueryResponse<Q>(
+  liquidityActionDescription: string,
+  queryResponse?: Q
+): Q {
+  if (!queryResponse) {
+    // This should never happen because we don't allow the user to trigger buildLiquidityCallData
+    // before the query is loaded.
+    console.error(`Missing queryResponse in ${liquidityActionDescription}`)
+    throw new SentryError(
+      `Missing queryResponse.
+It looks that you tried to call useBuildCallData before the last query finished generating queryResponse`
+    )
+  }
+
+  return queryResponse
+}

lib/modules/pool/actions/add-liquidity/add-liquidity.types.ts

@@ -5,7 +5,7 @@ export type QueryAddLiquidityOutput = {
   bptOut: TokenAmount
 }
 
-export type BuildAddLiquidityInputs = {
+export type BuildAddLiquidityInput = {
   account: Address
   slippagePercent: string
 }

lib/modules/pool/actions/add-liquidity/handlers/AddLiquidity.handler.ts

@@ -1,6 +1,6 @@
 import { TransactionConfig } from '@/lib/modules/web3/contracts/contract.types'
 import { HumanAmountIn } from '../../liquidity-types'
-import { BuildAddLiquidityInputs, QueryAddLiquidityOutput } from '../add-liquidity.types'
+import { BuildAddLiquidityInput, QueryAddLiquidityOutput } from '../add-liquidity.types'
 
 /**
  * AddLiquidityHandler is an interface that defines the methods that must be implemented by a handler.
@@ -24,7 +24,6 @@ export interface AddLiquidityHandler {
   /*
     Build tx callData payload for adding liquidity
     It is responsibility of the UI to avoid calling buildAddLiquidityCallData before the last queryAddLiquidity was finished
-    ("Next" button will be disabled if the react-query calling queryAddLiquidity is in loading state)
   */
-  buildAddLiquidityCallData(inputs: BuildAddLiquidityInputs): Promise<TransactionConfig>
+  buildAddLiquidityCallData(inputs: BuildAddLiquidityInput): Promise<TransactionConfig>
 }

lib/modules/pool/actions/add-liquidity/handlers/TwammAddLiquidity.handler.ts

@@ -3,7 +3,7 @@ import { TransactionConfig } from '@/lib/modules/web3/contracts/contract.types'
 import { emptyAddress } from '@/lib/modules/web3/contracts/wagmi-helpers'
 import { Token, TokenAmount } from '@balancer/sdk'
 import { HumanAmountIn } from '../../liquidity-types'
-import { BuildAddLiquidityInputs } from '../add-liquidity.types'
+import { BuildAddLiquidityInput } from '../add-liquidity.types'
 import { AddLiquidityHandler } from './AddLiquidity.handler'
 
 /**
@@ -35,7 +35,7 @@ export class TwammAddLiquidityHandler implements AddLiquidityHandler {
   // TODO: This is a non-sense example implementation
   public async buildAddLiquidityCallData({
     account,
-  }: BuildAddLiquidityInputs): Promise<TransactionConfig> {
+  }: BuildAddLiquidityInput): Promise<TransactionConfig> {
     if (!this.humanAmountsIn) {
       throw new Error(
         `Missing humanAmountsIn.

lib/modules/pool/actions/add-liquidity/handlers/UnbalancedAddLiquidity.handler.ts

@@ -11,11 +11,14 @@ import {
   Slippage,
 } from '@balancer/sdk'
 import { Pool } from '../../../usePool'
-import { LiquidityActionHelpers, areEmptyAmounts } from '../../LiquidityActionHelpers'
+import {
+  LiquidityActionHelpers,
+  areEmptyAmounts,
+  ensureLastQueryResponse,
+} from '../../LiquidityActionHelpers'
 import { HumanAmountIn } from '../../liquidity-types'
-import { BuildAddLiquidityInputs, QueryAddLiquidityOutput } from '../add-liquidity.types'
+import { BuildAddLiquidityInput, QueryAddLiquidityOutput } from '../add-liquidity.types'
 import { AddLiquidityHandler } from './AddLiquidity.handler'
-import { SentryError } from '@/lib/shared/utils/errors'
 
 /**
  * UnbalancedAddLiquidityHandler is a handler that implements the
@@ -61,22 +64,11 @@ export class UnbalancedAddLiquidityHandler implements AddLiquidityHandler {
     return priceImpactABA.decimal
   }
 
-  /*
-    sdkQueryOutput is the result of the query that we run in the add liquidity form
-  */
   public async buildAddLiquidityCallData({
     account,
     slippagePercent,
-  }: BuildAddLiquidityInputs): Promise<TransactionConfig> {
-    if (!this.queryResponse) {
-      // This should never happen because we don't allow the user to trigger buildAddLiquidityCallData
-      // before the query is loaded.
-      console.error('Missing queryResponse.')
-      throw new SentryError(
-        `Missing queryResponse.
-It looks that you tried to call useBuildCallData before the last query finished generating queryResponse`
-      )
-    }
+  }: BuildAddLiquidityInput): Promise<TransactionConfig> {
+    this.queryResponse = ensureLastQueryResponse('Unbalanced add liquidity', this.queryResponse)
 
     const addLiquidity = new AddLiquidity()
 

lib/modules/pool/actions/remove-liquidity/handlers/ProportionalRemoveLiquidity.handler.integration.spec.ts

@@ -4,26 +4,30 @@ import { aBalWethPoolElementMock } from '@/test/msw/builders/gqlPoolElement.buil
 import { defaultTestUserAccount } from '@/test/utils/wagmi'
 import { aPhantomStablePoolStateInputMock } from '../../../__mocks__/pool.builders'
 import { Pool } from '../../../usePool'
-import { RemoveLiquidityInputs, RemoveLiquidityType } from '../remove-liquidity.types'
+import { QueryRemoveLiquidityInput, RemoveLiquidityType } from '../remove-liquidity.types'
 import { selectRemoveLiquidityHandler } from './selectRemoveLiquidityHandler'
+import { ProportionalRemoveLiquidityHandler } from './ProportionalRemoveLiquidity.handler'
 
 const poolMock = aBalWethPoolElementMock() // 80BAL-20WETH
 
-function selectProportionalHandler(pool: Pool) {
-  return selectRemoveLiquidityHandler(pool, RemoveLiquidityType.Proportional)
+function selectProportionalHandler(pool: Pool): ProportionalRemoveLiquidityHandler {
+  return selectRemoveLiquidityHandler(
+    pool,
+    RemoveLiquidityType.Proportional
+  ) as ProportionalRemoveLiquidityHandler
 }
 
-const inputs: RemoveLiquidityInputs = {
+const defaultQueryInput: QueryRemoveLiquidityInput = {
   humanBptIn: '1',
-  account: defaultTestUserAccount,
-  slippagePercent: '0.2',
 }
 
+const defaultBuildInput = { account: defaultTestUserAccount, slippagePercent: '0.2' }
+
 describe('When proportionally removing liquidity for a weighted pool', () => {
   test('returns ZERO price impact', async () => {
     const handler = selectProportionalHandler(poolMock)
 
-    const result = await handler.queryRemoveLiquidity(inputs)
+    const result = await handler.queryRemoveLiquidity(defaultQueryInput)
 
     const [balTokenAmountOut, wEthTokenAmountOut] = result.amountsOut
 
@@ -36,7 +40,7 @@ describe('When proportionally removing liquidity for a weighted pool', () => {
   test('queries amounts out', async () => {
     const handler = selectProportionalHandler(poolMock)
 
-    const result = await handler.queryRemoveLiquidity(inputs)
+    const result = await handler.queryRemoveLiquidity(defaultQueryInput)
 
     const [balTokenAmountOut, wEthTokenAmountOut] = result.amountsOut
 
@@ -50,9 +54,9 @@ describe('When proportionally removing liquidity for a weighted pool', () => {
   test('builds Tx Config', async () => {
     const handler = selectProportionalHandler(poolMock)
 
-    const { sdkQueryOutput } = await handler.queryRemoveLiquidity(inputs)
+    await handler.queryRemoveLiquidity(defaultQueryInput)
 
-    const result = await handler.buildRemoveLiquidityTx({ inputs, sdkQueryOutput })
+    const result = await handler.buildRemoveLiquidityCallData(defaultBuildInput)
 
     expect(result.to).toBe(networkConfig.contracts.balancer.vaultV2)
     expect(result.data).toBeDefined()
@@ -65,9 +69,9 @@ describe('When removing liquidity from a stable pool', () => {
 
     const handler = selectProportionalHandler(pool)
 
-    const { sdkQueryOutput } = await handler.queryRemoveLiquidity(inputs)
+    await handler.queryRemoveLiquidity(defaultQueryInput)
 
-    const result = await handler.buildRemoveLiquidityTx({ inputs, sdkQueryOutput })
+    const result = await handler.buildRemoveLiquidityCallData(defaultBuildInput)
     expect(result.account).toBe(defaultTestUserAccount)
   })
 })

lib/modules/pool/actions/remove-liquidity/handlers/ProportionalRemoveLiquidity.handler.ts

@@ -12,61 +12,56 @@ import {
 import { Address, parseEther } from 'viem'
 import { BPT_DECIMALS } from '../../../pool.constants'
 import { Pool } from '../../../usePool'
-import { LiquidityActionHelpers } from '../../LiquidityActionHelpers'
+import { LiquidityActionHelpers, ensureLastQueryResponse } from '../../LiquidityActionHelpers'
 import {
-  BuildLiquidityInputs,
-  RemoveLiquidityInputs,
-  RemoveLiquidityOutputs,
+  BuildRemoveLiquidityInput,
+  QueryRemoveLiquidityInput,
+  QueryRemoveLiquidityOutput,
 } from '../remove-liquidity.types'
 import { RemoveLiquidityHandler } from './RemoveLiquidity.handler'
 
 export class ProportionalRemoveLiquidityHandler implements RemoveLiquidityHandler {
   helpers: LiquidityActionHelpers
-  sdkQueryOutput?: RemoveLiquidityQueryOutput
+  queryResponse?: RemoveLiquidityQueryOutput
 
   constructor(pool: Pool) {
     this.helpers = new LiquidityActionHelpers(pool)
   }
 
   public async queryRemoveLiquidity({
     humanBptIn: bptIn,
-  }: RemoveLiquidityInputs): Promise<RemoveLiquidityOutputs> {
+  }: QueryRemoveLiquidityInput): Promise<QueryRemoveLiquidityOutput> {
+    // Deletes the previous queryResponse to enforce that we don't buildCallData with an outdated queryResponse (while a new one is loading)
+    this.queryResponse = undefined
     const removeLiquidity = new RemoveLiquidity()
     const removeLiquidityInput = this.constructSdkInput(bptIn)
 
-    this.sdkQueryOutput = await removeLiquidity.query(
+    this.queryResponse = await removeLiquidity.query(
       removeLiquidityInput,
       this.helpers.poolStateInput
     )
 
-    return { amountsOut: this.sdkQueryOutput.amountsOut }
+    return { amountsOut: this.queryResponse.amountsOut }
   }
 
   public async calculatePriceImpact(): Promise<number> {
     // proportional remove liquidity does not have price impact
     return 0
   }
 
-  /*
-    sdkQueryOutput is the result of the query that we run in the remove liquidity form
-  */
-  public async buildRemoveLiquidityTx(
-    buildInputs: BuildLiquidityInputs
-  ): Promise<TransactionConfig> {
-    const { account, slippagePercent } = buildInputs.inputs
-    if (!account || !slippagePercent) throw new Error('Missing account or slippage')
-    if (!this.sdkQueryOutput) {
-      console.error('Missing sdkQueryOutput in buildRemoveLiquidityTx')
-      throw new Error(
-        `Missing sdkQueryOutput.
-It looks that you did not call useRemoveLiquidityBtpOutQuery before trying to build the tx config`
-      )
-    }
+  public async buildRemoveLiquidityCallData({
+    account,
+    slippagePercent,
+  }: BuildRemoveLiquidityInput): Promise<TransactionConfig> {
+    this.queryResponse = ensureLastQueryResponse(
+      'Proportional remove liquidity',
+      this.queryResponse
+    )
 
     const removeLiquidity = new RemoveLiquidity()
 
     const { call, to, value } = removeLiquidity.buildCall({
-      ...this.sdkQueryOutput,
+      ...this.queryResponse,
       slippage: Slippage.fromPercentage(`${Number(slippagePercent)}`),
       sender: account,
       recipient: account,

lib/modules/pool/actions/remove-liquidity/handlers/RemoveLiquidity.handler.ts

@@ -1,8 +1,8 @@
 import { TransactionConfig } from '@/lib/modules/web3/contracts/contract.types'
 import {
-  RemoveLiquidityInputs,
-  RemoveLiquidityOutputs,
-  BuildLiquidityInputs,
+  QueryRemoveLiquidityInput,
+  BuildRemoveLiquidityInput,
+  QueryRemoveLiquidityOutput,
 } from '../remove-liquidity.types'
 
 /**
@@ -12,9 +12,12 @@ import {
  */
 export interface RemoveLiquidityHandler {
   // Query the SDK for the expected output of removing liquidity
-  queryRemoveLiquidity(inputs: RemoveLiquidityInputs): Promise<RemoveLiquidityOutputs>
+  queryRemoveLiquidity(inputs: QueryRemoveLiquidityInput): Promise<QueryRemoveLiquidityOutput>
   // Calculate the price impact of removing liquidity
-  calculatePriceImpact(inputs: RemoveLiquidityInputs): Promise<number>
-  // Build tx payload for removing liquidity
-  buildRemoveLiquidityTx(inputs: BuildLiquidityInputs): Promise<TransactionConfig>
+  calculatePriceImpact(inputs: QueryRemoveLiquidityInput): Promise<number>
+  /*
+    Build tx callData payload for removing liquidity
+    It is responsibility of the UI to avoid calling buildRemoveLiquidityCallData before the last queryRemoveLiquidity was finished
+  */
+  buildRemoveLiquidityCallData(inputs: BuildRemoveLiquidityInput): Promise<TransactionConfig>
 }

lib/modules/pool/actions/remove-liquidity/handlers/SingleTokenRemoveLiquidity.handler.integration.spec.ts

@@ -3,26 +3,31 @@ import { balAddress, wETHAddress } from '@/lib/debug-helpers'
 import { aBalWethPoolElementMock } from '@/test/msw/builders/gqlPoolElement.builders'
 import { defaultTestUserAccount } from '@/test/utils/wagmi'
 import { Pool } from '../../../usePool'
-import { RemoveLiquidityInputs, RemoveLiquidityType } from '../remove-liquidity.types'
+import { QueryRemoveLiquidityInput, RemoveLiquidityType } from '../remove-liquidity.types'
 import { selectRemoveLiquidityHandler } from './selectRemoveLiquidityHandler'
+import { SingleTokenRemoveLiquidityHandler } from './SingleTokenRemoveLiquidity.handler'
 
 const poolMock = aBalWethPoolElementMock() // 80BAL-20WETH
 
-function selectSingleTokenHandler(pool: Pool) {
-  return selectRemoveLiquidityHandler(pool, RemoveLiquidityType.SingleToken)
+function selectSingleTokenHandler(pool: Pool): SingleTokenRemoveLiquidityHandler {
+  return selectRemoveLiquidityHandler(
+    pool,
+    RemoveLiquidityType.SingleToken
+  ) as SingleTokenRemoveLiquidityHandler
 }
 
+const defaultQueryInput: QueryRemoveLiquidityInput = {
+  humanBptIn: '1',
+  tokenOut: balAddress,
+}
+
+const defaultBuildInput = { account: defaultTestUserAccount, slippagePercent: '0.2' }
+
 describe('When removing unbalanced liquidity for a weighted pool', () => {
   test('queries amounts out', async () => {
-    const inputs: RemoveLiquidityInputs = {
-      humanBptIn: '1',
-      account: defaultTestUserAccount,
-      tokenOut: balAddress,
-    }
-
     const handler = selectSingleTokenHandler(poolMock)
 
-    const result = await handler.queryRemoveLiquidity(inputs)
+    const result = await handler.queryRemoveLiquidity(defaultQueryInput)
 
     const [balTokenAmountOut, wEthTokenAmountOut] = result.amountsOut
 
@@ -36,16 +41,14 @@ describe('When removing unbalanced liquidity for a weighted pool', () => {
   test('builds Tx Config', async () => {
     const handler = selectSingleTokenHandler(poolMock)
 
-    const inputs: RemoveLiquidityInputs = {
+    const inputs: QueryRemoveLiquidityInput = {
       humanBptIn: '1',
-      account: defaultTestUserAccount,
-      slippagePercent: '0.2',
       tokenOut: balAddress,
     }
 
-    const { sdkQueryOutput } = await handler.queryRemoveLiquidity(inputs)
+    await handler.queryRemoveLiquidity(inputs)
 
-    const result = await handler.buildRemoveLiquidityTx({ inputs, sdkQueryOutput })
+    const result = await handler.buildRemoveLiquidityCallData(defaultBuildInput)
 
     expect(result.to).toBe(networkConfig.contracts.balancer.vaultV2)
     expect(result.data).toBeDefined()