fix: Zap fixes

pactsdk/add_liquidity.py

@@ -5,13 +5,15 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, cast
 
+from pactsdk.exceptions import PactSdkError
 from pactsdk.stableswap_calculator import (
     StableswapCalculator,
     StableswapParams,
     get_add_liquidity_bonus_pct,
     get_tx_fee,
 )
 
+from .constant_product_calculator import get_constant_product_minted_liquidity_tokens
 from .transaction_group import TransactionGroup
 
 if TYPE_CHECKING:
@@ -87,11 +89,6 @@ def build_effect(self) -> AddLiquidityEffect:
         swap_calc = self.pool.calculator.swap_calculator
         state = self.pool.state
 
-        minted_liquidity_tokens = swap_calc.get_minted_liquidity_tokens(
-            self.primary_asset_amount,
-            self.secondary_asset_amount,
-        )
-
         if isinstance(swap_calc, StableswapCalculator):
             i_amplifier = swap_calc.get_amplifier()
             params = cast(StableswapParams, self.pool.params)
@@ -104,10 +101,26 @@ def build_effect(self) -> AddLiquidityEffect:
                 i_amplifier,
                 params.precision,
             )
+            minted_liquidity_tokens = swap_calc.get_minted_liquidity_tokens(
+                self.primary_asset_amount,
+                self.secondary_asset_amount,
+            )
             amplifier = i_amplifier / (self.pool.internal_state.PRECISION or 1)
 
             # 1 for each invariant calculation (3) and 1 for sending liquidity tokens.
             tx_fee = get_tx_fee(swap_calc.mint_tokens_invariant_iterations, 4)
+        else:
+            minted_liquidity_tokens = get_constant_product_minted_liquidity_tokens(
+                self.primary_asset_amount,
+                self.secondary_asset_amount,
+                state.total_primary,
+                state.total_secondary,
+                state.total_liquidity,
+            )
+            if minted_liquidity_tokens <= 0:
+                raise PactSdkError(
+                    "Amount of minted liquidity tokens must be greater then 0.",
+                )
 
         return AddLiquidityEffect(
             minted_liquidity_tokens=minted_liquidity_tokens,

pactsdk/pool.py

@@ -344,7 +344,7 @@ def build_add_liquidity_txs(
                         secondary_asset_amount=secondary_small_amount,
                         suggested_params=suggested_params,
                         fee=liquidity_addition.effect.tx_fee,
-                        note=b"Initial add liquidity",
+                        note=b"Pact initial add liquidity",
                     )
 
         txs = self.build_raw_add_liquidity_txs(
@@ -370,12 +370,14 @@ def build_raw_add_liquidity_txs(
             address=address,
             asset=self.primary_asset,
             amount=primary_asset_amount,
+            note=b"Pact add liquidity deposit",
             suggested_params=suggested_params,
         )
         tx2 = self._make_deposit_tx(
             address=address,
             asset=self.secondary_asset,
             amount=secondary_asset_amount,
+            note=b"Pact add liquidity deposit",
             suggested_params=suggested_params,
         )
         tx3 = self._make_application_noop_tx(
@@ -425,6 +427,7 @@ def build_remove_liquidity_txs(
             address=address,
             amount=amount,
             asset=self.liquidity_asset,
+            note=b"Pact remove liquidity deposit",
             suggested_params=suggested_params,
         )
         tx2 = self._make_application_noop_tx(
@@ -493,6 +496,7 @@ def build_swap_txs(
             address=address,
             amount=swap.effect.amount_deposited,
             asset=swap.asset_deposited,
+            note=b"Pact swap deposit",
             suggested_params=suggested_params,
         )
         tx2 = self._make_application_noop_tx(
@@ -593,6 +597,7 @@ def _make_deposit_tx(
         asset: Asset,
         address: str,
         amount: int,
+        note: bytes,
         suggested_params: transaction.SuggestedParams,
     ):
         if not asset.index:
@@ -601,13 +606,15 @@ def _make_deposit_tx(
                 sender=address,
                 receiver=self.get_escrow_address(),
                 amt=amount,
+                note=note,
                 sp=suggested_params,
             )
 
         return transaction.AssetTransferTxn(
             sender=address,
             receiver=self.get_escrow_address(),
             amt=amount,
+            note=note,
             sp=suggested_params,
             index=asset.index,
         )

pactsdk/zap.py

@@ -1,5 +1,6 @@
 """Set of utility classes for managing and performing zaps.
 """
+import copy
 import math
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
@@ -119,11 +120,7 @@ def __post_init__(self):
         self.swap = self.pool.prepare_swap(
             self.asset, self.params.swap_deposited, self.slippage_pct
         )
-        self.liquidity_addition = LiquidityAddition(
-            pool=self.pool,
-            primary_asset_amount=self.params.primary_add_liq,
-            secondary_asset_amount=self.params.secondary_add_liq,
-        )
+        self.liquidity_addition = self.prepare_add_liq()
 
     def prepare_tx_group(self, address: str) -> TransactionGroup:
         """Creates the transactions needed to perform zap and returns them as a transaction group ready to be signed and committed.
@@ -149,7 +146,26 @@ def get_zap_params(self) -> ZapParams:
         )
         if not self._is_asset_primary():
             temp = params.primary_add_liq
-            params.primary_add_liq = params.secondary_add_liq
+            params.primary_add_liq = params.secondary_add_liq - 1
             params.secondary_add_liq = temp
+        else:
+            params.secondary_add_liq -= 1
 
         return params
+
+    def prepare_add_liq(self):
+        updated_state = self.pool.state
+        if self._is_asset_primary():
+            updated_state.total_primary += self.params.swap_deposited
+            updated_state.total_secondary -= self.params.secondary_add_liq
+        else:
+            updated_state.total_primary -= self.params.primary_add_liq
+            updated_state.total_secondary += self.params.swap_deposited
+
+        pool = copy.copy(self.pool)
+        pool.state = updated_state
+        return LiquidityAddition(
+            pool=pool,
+            primary_asset_amount=self.params.primary_add_liq,
+            secondary_asset_amount=self.params.secondary_add_liq,
+        )

tests/test_zap.py

@@ -23,7 +23,7 @@ def test_calculate_zap_params():
     assert zap_primary_add.params == pactsdk.ZapParams(
         swap_deposited=4888,
         primary_add_liq=5112,
-        secondary_add_liq=4646,
+        secondary_add_liq=4645,
     )
 
     #  Perform a zap using secondary asset.
@@ -32,7 +32,7 @@ def test_calculate_zap_params():
     )
     assert zap_secondary_add.params == pactsdk.ZapParams(
         swap_deposited=4888,
-        primary_add_liq=4646,
+        primary_add_liq=4645,
         secondary_add_liq=5112,
     )
 
@@ -44,7 +44,7 @@ def test_calculate_zap_params():
     ubalanced_zap = testbed2.pool.prepare_zap(testbed2.pool.secondary_asset, 20_000, 2)
     assert ubalanced_zap.params == pactsdk.ZapParams(
         swap_deposited=7339,
-        primary_add_liq=42199,
+        primary_add_liq=42198,
         secondary_add_liq=12661,
     )
 
@@ -102,9 +102,9 @@ def test_zap_e2e():
     testbed.pool.update_state()
 
     assert testbed.pool.state == pactsdk.PoolState(
-        total_liquidity=104872,
-        total_primary=109999,
-        total_secondary=100000,
-        primary_asset_price=0.9090991736288512,
-        secondary_asset_price=1.09999,
+        total_liquidity=104871,
+        total_primary=109998,
+        total_secondary=99999,
+        primary_asset_price=0.9090983472426772,
+        secondary_asset_price=1.099990999909999,
     )