Add common_test_functions module and update add_remove_lp() test to use it.

frame/pablo/src/common_test_functions.rs

@@ -0,0 +1,66 @@
+use crate::{
+	mock::{Pablo, *},
+	PoolConfiguration::{ConstantProduct, StableSwap},
+	PoolInitConfiguration,
+};
+use frame_support::{
+	assert_ok,
+	traits::fungibles::{Inspect, Mutate},
+};
+
+/// `expected_lp_check` takes base_amount, quote_amount and lp_tokens in order and returns
+/// true if lp_tokens are expected for given base_amount, quote_amount.
+pub fn common_add_remove_lp(
+	init_config: PoolInitConfiguration<AssetId>,
+	init_base_amount: Balance,
+	init_quote_amount: Balance,
+	base_amount: Balance,
+	quote_amount: Balance,
+	expected_lp_check: impl Fn(Balance, Balance, Balance) -> bool,
+) {
+	let pool_id = Pablo::do_create_pool(&ALICE, init_config.clone()).expect("pool creation failed");
+	let pair = match init_config {
+		PoolInitConfiguration::StableSwap { pair, .. } => pair,
+		PoolInitConfiguration::ConstantProduct { pair, .. } => pair,
+	};
+	// Mint the tokens
+	assert_ok!(Tokens::mint_into(pair.base, &ALICE, init_base_amount));
+	assert_ok!(Tokens::mint_into(pair.quote, &ALICE, init_quote_amount));
+
+	// Add the liquidity
+	assert_ok!(Pablo::add_liquidity(
+		Origin::signed(ALICE),
+		pool_id,
+		init_base_amount,
+		init_quote_amount,
+		0,
+		false
+	));
+
+	let pool = Pablo::pools(pool_id).expect("pool not found");
+	let lp_token = match pool {
+		StableSwap(pool) => pool.lp_token,
+		ConstantProduct(pool) => pool.lp_token,
+	};
+	// Mint the tokens
+	assert_ok!(Tokens::mint_into(pair.base, &BOB, base_amount));
+	assert_ok!(Tokens::mint_into(pair.quote, &BOB, quote_amount));
+
+	let lp = Tokens::balance(lp_token, &BOB);
+	assert_eq!(lp, 0_u128);
+	// Add the liquidity
+	assert_ok!(Pablo::add_liquidity(
+		Origin::signed(BOB),
+		pool_id,
+		base_amount,
+		quote_amount,
+		0,
+		false
+	));
+	let lp = Tokens::balance(lp_token, &BOB);
+	assert!(expected_lp_check(base_amount, quote_amount, lp));
+	assert_ok!(Pablo::remove_liquidity(Origin::signed(BOB), pool_id, lp, 0, 0));
+	let lp = Tokens::balance(lp_token, &BOB);
+	// all lp tokens must have been burnt
+	assert_eq!(lp, 0_u128);
+}

frame/pablo/src/lib.rs

@@ -35,6 +35,8 @@
 
 pub use pallet::*;
 
+#[cfg(test)]
+mod common_test_functions;
 #[cfg(test)]
 mod mock;
 #[cfg(test)]
@@ -417,32 +419,6 @@ pub mod pallet {
 		pub(crate) fn account_id(pool_id: &T::PoolId) -> T::AccountId {
 			T::PalletId::get().into_sub_account(pool_id)
 		}
-
-		fn do_compute_swap(
-			pool_id: T::PoolId,
-			pair: CurrencyPair<T::AssetId>,
-			quote_amount: T::Balance,
-			apply_fees: bool,
-			fee: Permill,
-			protocol_fee: Permill,
-		) -> Result<(T::Balance, T::Balance, T::Balance, T::Balance), DispatchError> {
-			let base_amount = Self::get_exchange_value(pool_id, pair.base, quote_amount)?;
-			let base_amount_u: u128 = T::Convert::convert(base_amount);
-
-			let (lp_fee, protocol_fee) = if apply_fees {
-				let lp_fee = fee.mul_floor(base_amount_u);
-				// protocol_fee is computed based on lp_fee
-				let protocol_fee = protocol_fee.mul_floor(lp_fee);
-				let lp_fee = T::Convert::convert(lp_fee);
-				let protocol_fee = T::Convert::convert(protocol_fee);
-				(lp_fee, protocol_fee)
-			} else {
-				(T::Balance::zero(), T::Balance::zero())
-			};
-
-			let base_amount_excluding_fees = base_amount.safe_sub(&lp_fee)?;
-			Ok((base_amount_excluding_fees, quote_amount, lp_fee, protocol_fee))
-		}
 	}
 
 	impl<T: Config> Amm for Pallet<T> {
@@ -476,14 +452,14 @@ pub mod pallet {
 			match pool {
 				PoolConfiguration::StableSwap(stable_swap_pool_info) =>
 					StableSwap::<T>::get_exchange_value(
-						stable_swap_pool_info,
-						pool_account,
+						&stable_swap_pool_info,
+						&pool_account,
 						asset_id,
 						amount,
 					),
 				ConstantProduct(constant_product_pool_info) => Uniswap::<T>::get_exchange_value(
-					constant_product_pool_info,
-					pool_account,
+					&constant_product_pool_info,
+					&pool_account,
 					asset_id,
 					amount,
 				),
@@ -611,14 +587,7 @@ pub mod pallet {
 					// provided pair might have been swapped
 					ensure!(pair == pool.pair, Error::<T>::PairMismatch);
 					let (base_amount_excluding_fees, quote_amount, lp_fees, protocol_fees) =
-						Self::do_compute_swap(
-							pool_id,
-							pair,
-							quote_amount,
-							true,
-							pool.fee,
-							pool.protocol_fee,
-						)?;
+						StableSwap::<T>::do_compute_swap(&pool, &pool_account, quote_amount, true)?;
 
 					ensure!(
 						base_amount_excluding_fees >= min_receive,
@@ -657,7 +626,7 @@ pub mod pallet {
 					let (base_amount, quote_amount_excluding_fees, lp_fees, owner_fees) =
 						Uniswap::<T>::do_compute_swap(
 							&constant_product_pool_info,
-							pool_account,
+							&pool_account,
 							pair,
 							quote_amount,
 							true,
@@ -668,7 +637,6 @@ pub mod pallet {
 
 					ensure!(base_amount >= min_receive, Error::<T>::CannotRespectMinimumRequested);
 
-					let pool_account = Self::account_id(&pool_id);
 					T::Assets::transfer(
 						pair.quote,
 						who,

frame/pablo/src/stable_swap.rs

@@ -74,14 +74,14 @@ impl<T: Config> StableSwap<T> {
 	}
 
 	pub fn get_exchange_value(
-		pool: StableSwapPoolInfo<T::AccountId, T::AssetId>,
-		pool_account: T::AccountId,
+		pool: &StableSwapPoolInfo<T::AccountId, T::AssetId>,
+		pool_account: &T::AccountId,
 		asset_id: T::AssetId,
 		amount: T::Balance,
 	) -> Result<T::Balance, DispatchError> {
 		let pair = if asset_id == pool.pair.base { pool.pair } else { pool.pair.swap() };
-		let pool_base_aum = T::Assets::balance(pair.base, &pool_account);
-		let pool_quote_aum = T::Assets::balance(pair.quote, &pool_account);
+		let pool_base_aum = T::Assets::balance(pair.base, pool_account);
+		let pool_quote_aum = T::Assets::balance(pair.quote, pool_account);
 		let amp = T::Convert::convert(pool.amplification_coefficient.into());
 		let d = Self::get_invariant(pool_base_aum, pool_quote_aum, amp)?;
 		let new_quote_amount = pool_quote_aum.safe_add(&amount)?;
@@ -106,6 +106,31 @@ impl<T: Config> StableSwap<T> {
 		Ok(difference)
 	}
 
+	pub fn do_compute_swap(
+		pool: &StableSwapPoolInfo<T::AccountId, T::AssetId>,
+		pool_account: &T::AccountId,
+		quote_amount: T::Balance,
+		apply_fees: bool,
+	) -> Result<(T::Balance, T::Balance, T::Balance, T::Balance), DispatchError> {
+		let base_amount =
+			Self::get_exchange_value(pool, pool_account, pool.pair.base, quote_amount)?;
+		let base_amount_u: u128 = T::Convert::convert(base_amount);
+
+		let (lp_fee, protocol_fee) = if apply_fees {
+			let lp_fee = pool.fee.mul_floor(base_amount_u);
+			// protocol_fee is computed based on lp_fee
+			let protocol_fee = pool.protocol_fee.mul_floor(lp_fee);
+			let lp_fee = T::Convert::convert(lp_fee);
+			let protocol_fee = T::Convert::convert(protocol_fee);
+			(lp_fee, protocol_fee)
+		} else {
+			(T::Balance::zero(), T::Balance::zero())
+		};
+
+		let base_amount_excluding_fees = base_amount.safe_sub(&lp_fee)?;
+		Ok((base_amount_excluding_fees, quote_amount, lp_fee, protocol_fee))
+	}
+
 	pub fn add_liquidity(
 		who: &T::AccountId,
 		pool: StableSwapPoolInfo<T::AccountId, T::AssetId>,

frame/pablo/src/stable_swap_tests.rs

@@ -1,6 +1,7 @@
 #![allow(unreachable_patterns)] // TODO: remove this when there are more pool configurations added.
 
 use crate::{
+	common_test_functions::*,
 	mock::{Pablo, *},
 	PoolConfiguration::StableSwap,
 	PoolInitConfiguration,
@@ -147,47 +148,29 @@ fn test_dex_demo() {
 #[test]
 fn add_remove_lp() {
 	new_test_ext().execute_with(|| {
+		let pool_init_config = PoolInitConfiguration::StableSwap {
+			pair: CurrencyPair::new(USDC, USDT),
+			amplification_coefficient: 10_u16,
+			fee: Permill::zero(),
+			protocol_fee: Permill::zero(),
+		};
 		let unit = 1_000_000_000_000_u128;
 		let initial_usdt = 1_000_000_000_000_u128 * unit;
 		let initial_usdc = 1_000_000_000_000_u128 * unit;
-		let pool_id = create_stable_swap_pool(
-			USDC,
-			USDT,
+		let usdc_amount = 1000 * unit;
+		let usdt_amount = 1000 * unit;
+		let expected_lp_check = |base_amount: Balance,
+		                         quote_amount: Balance,
+		                         lp: Balance|
+		 -> bool { base_amount + quote_amount == lp };
+		common_add_remove_lp(
+			pool_init_config,
 			initial_usdc,
 			initial_usdt,
-			100_u16,
-			Permill::zero(),
-			Permill::zero(),
+			usdc_amount,
+			usdt_amount,
+			expected_lp_check,
 		);
-		let pool = Pablo::pools(pool_id).expect("pool not found");
-		let pool = match pool {
-			StableSwap(pool) => pool,
-			_ => panic!("expected stable_swap pool"),
-		};
-		let bob_usdc = 1000 * unit;
-		let bob_usdt = 1000 * unit;
-		// Mint the tokens
-		assert_ok!(Tokens::mint_into(USDC, &BOB, bob_usdc));
-		assert_ok!(Tokens::mint_into(USDT, &BOB, bob_usdt));
-
-		let lp = Tokens::balance(pool.lp_token, &BOB);
-		assert_eq!(lp, 0_u128);
-		// Add the liquidity
-		assert_ok!(Pablo::add_liquidity(
-			Origin::signed(BOB),
-			pool_id,
-			bob_usdc,
-			bob_usdt,
-			0,
-			false
-		));
-		let lp = Tokens::balance(pool.lp_token, &BOB);
-		// must have received some lp tokens
-		assert!(lp == bob_usdt + bob_usdc);
-		assert_ok!(Pablo::remove_liquidity(Origin::signed(BOB), pool_id, lp, 0, 0));
-		let lp = Tokens::balance(pool.lp_token, &BOB);
-		// all lp tokens must have been burnt
-		assert_eq!(lp, 0_u128);
 	});
 }
 

frame/pablo/src/uniswap.rs

@@ -61,16 +61,15 @@ impl<T: Config> Uniswap<T> {
 	}
 
 	pub(crate) fn get_exchange_value(
-		pool: ConstantProductPoolInfo<T::AccountId, T::AssetId>,
-		pool_account: T::AccountId,
+		pool: &ConstantProductPoolInfo<T::AccountId, T::AssetId>,
+		pool_account: &T::AccountId,
 		asset_id: T::AssetId,
 		amount: T::Balance,
 	) -> Result<T::Balance, DispatchError> {
 		let amount = T::Convert::convert(amount);
 		let half_weight = Permill::from_percent(50);
-		let pool_base_aum = T::Convert::convert(T::Assets::balance(pool.pair.base, &pool_account));
-		let pool_quote_aum =
-			T::Convert::convert(T::Assets::balance(pool.pair.quote, &pool_account));
+		let pool_base_aum = T::Convert::convert(T::Assets::balance(pool.pair.base, pool_account));
+		let pool_quote_aum = T::Convert::convert(T::Assets::balance(pool.pair.quote, pool_account));
 		let exchange_amount = if asset_id == pool.pair.quote {
 			compute_out_given_in(half_weight, half_weight, pool_quote_aum, pool_base_aum, amount)
 		} else {
@@ -154,13 +153,13 @@ impl<T: Config> Uniswap<T> {
 
 	pub(crate) fn do_compute_swap(
 		pool: &ConstantProductPoolInfo<T::AccountId, T::AssetId>,
-		pool_account: T::AccountId,
+		pool_account: &T::AccountId,
 		pair: CurrencyPair<T::AssetId>,
 		quote_amount: T::Balance,
 		apply_fees: bool,
 	) -> Result<(T::Balance, T::Balance, T::Balance, T::Balance), DispatchError> {
-		let pool_base_aum = T::Convert::convert(T::Assets::balance(pair.base, &pool_account));
-		let pool_quote_aum = T::Convert::convert(T::Assets::balance(pair.quote, &pool_account));
+		let pool_base_aum = T::Convert::convert(T::Assets::balance(pair.base, pool_account));
+		let pool_quote_aum = T::Convert::convert(T::Assets::balance(pair.quote, pool_account));
 		let quote_amount = T::Convert::convert(quote_amount);
 
 		// https://uniswap.org/whitepaper.pdf

frame/pablo/src/uniswap_tests.rs

@@ -1,4 +1,5 @@
 use crate::{
+	common_test_functions::*,
 	mock::{Pablo, *},
 	PoolConfiguration::ConstantProduct,
 	PoolInitConfiguration,
@@ -145,30 +146,27 @@ fn test() {
 #[test]
 fn add_remove_lp() {
 	new_test_ext().execute_with(|| {
+		let pool_init_config = PoolInitConfiguration::ConstantProduct {
+			pair: CurrencyPair::new(BTC, USDT),
+			fee: Permill::zero(),
+			owner_fee: Permill::zero(),
+		};
 		let unit = 1_000_000_000_000_u128;
 		let initial_btc = 1_00_u128 * unit;
 		let btc_price = 45_000_u128;
 		let initial_usdt = initial_btc * btc_price;
-		let pool_id =
-			create_pool(BTC, USDT, initial_btc, initial_usdt, Permill::zero(), Permill::zero());
-		let pool = get_pool(pool_id);
-		let bob_btc = 10 * unit;
-		let bob_usdt = bob_btc * btc_price;
-		// Mint the tokens
-		assert_ok!(Tokens::mint_into(BTC, &BOB, bob_btc));
-		assert_ok!(Tokens::mint_into(USDT, &BOB, bob_usdt));
-
-		let lp = Tokens::balance(pool.lp_token, &BOB);
-		assert_eq!(lp, 0_u128);
-		// Add the liquidity
-		assert_ok!(<Pablo as Amm>::add_liquidity(&BOB, pool_id, bob_btc, bob_usdt, 0, false));
-		let lp = Tokens::balance(pool.lp_token, &BOB);
-		// must have received some lp tokens
-		assert!(lp > 0_u128);
-		assert_ok!(<Pablo as Amm>::remove_liquidity(&BOB, pool_id, lp, 0, 0));
-		let lp = Tokens::balance(pool.lp_token, &BOB);
-		// all lp tokens must have been burnt
-		assert_eq!(lp, 0_u128);
+		let btc_amount = 10 * unit;
+		let usdt_amount = btc_amount * btc_price;
+		let expected_lp_check =
+			|_base_amount: Balance, _quote_amount: Balance, lp: Balance| -> bool { lp > 0_u128 };
+		common_add_remove_lp(
+			pool_init_config,
+			initial_btc,
+			initial_usdt,
+			btc_amount,
+			usdt_amount,
+			expected_lp_check,
+		);
 	});
 }