From 3d38217f6870d94127d72c869cc3af6a89429838 Mon Sep 17 00:00:00 2001
From: Lars Lubkoll <11710767+lubkoll@users.noreply.github.com>
Date: Fri, 30 Aug 2024 10:26:44 +0000
Subject: [PATCH] Fix proptest

---
 .../cl-vault/tests/test-tube/proptest.rs      | 10 ++--
 .../cl-vault/tests/test-tube/setup.rs         | 55 +++++++++++++------
 2 files changed, 43 insertions(+), 22 deletions(-)

diff --git a/smart-contracts/osmosis/contracts/cl-vault/tests/test-tube/proptest.rs b/smart-contracts/osmosis/contracts/cl-vault/tests/test-tube/proptest.rs
index 667b283de..387757a29 100644
--- a/smart-contracts/osmosis/contracts/cl-vault/tests/test-tube/proptest.rs
+++ b/smart-contracts/osmosis/contracts/cl-vault/tests/test-tube/proptest.rs
@@ -1,6 +1,6 @@
 use crate::setup::{
-    get_event_attributes_by_ty_and_key, init_test_contract, MAX_SLIPPAGE_HIGH,
-    PERFORMANCE_FEE_DEFAULT,
+    get_event_attributes_by_ty_and_key, init_test_contract_with_dex_router_and_swap_pools,
+    MAX_SLIPPAGE_HIGH, PERFORMANCE_FEE_DEFAULT,
 };
 
 use cl_vault::{
@@ -335,8 +335,9 @@ proptest! {
         percentages in get_percentage_list(),
         account_indexes in get_account_index_list()
     ) {
-        let (app, contract_address, _cl_pool_id, admin_account, _) = init_test_contract(
+        let (app, contract_address, _dex_router_address, _cl_pool_id, _pools, admin_account, _) = init_test_contract_with_dex_router_and_swap_pools(
             "./test-tube-build/wasm32-unknown-unknown/release/cl_vault.wasm",
+            "../dex-router-osmosis/test-tube-build/wasm32-unknown-unknown/release/dex_router_osmosis.wasm",
             &[
                 Coin::new(340282366920938463463374607431768211455, "uosmo"),
                 Coin::new(340282366920938463463374607431768211455, DENOM_BASE),
@@ -363,7 +364,8 @@ proptest! {
             ],
             Uint128::zero(),
             Uint128::zero(),
-            PERFORMANCE_FEE_DEFAULT
+            PERFORMANCE_FEE_DEFAULT,
+            false
         );
         let wasm = Wasm::new(&app);
         let cl = ConcentratedLiquidity::new(&app);
diff --git a/smart-contracts/osmosis/contracts/cl-vault/tests/test-tube/setup.rs b/smart-contracts/osmosis/contracts/cl-vault/tests/test-tube/setup.rs
index fc7feb825..6921bf734 100644
--- a/smart-contracts/osmosis/contracts/cl-vault/tests/test-tube/setup.rs
+++ b/smart-contracts/osmosis/contracts/cl-vault/tests/test-tube/setup.rs
@@ -3,13 +3,18 @@
 use cl_vault::{helpers::generic::sort_tokens, msg::InstantiateMsg, state::VaultConfig};
 use cosmwasm_std::{coin, Addr, Attribute, Coin, Decimal, Uint128};
 use dex_router_osmosis::msg::{ExecuteMsg as DexExecuteMsg, InstantiateMsg as DexInstantiate};
-use osmosis_std::types::{
-    cosmos::{bank::v1beta1::QueryBalanceRequest, base::v1beta1},
-    cosmwasm::wasm::v1::MsgExecuteContractResponse,
-    osmosis::concentratedliquidity::v1beta1::{
-        CreateConcentratedLiquidityPoolsProposal, Pool, PoolRecord, PoolsRequest,
+use osmosis_std::{
+    try_proto_to_cosmwasm_coins,
+    types::{
+        cosmos::{bank::v1beta1::QueryBalanceRequest, base::v1beta1},
+        cosmwasm::wasm::v1::MsgExecuteContractResponse,
+        osmosis::{
+            concentratedliquidity::v1beta1::{
+                CreateConcentratedLiquidityPoolsProposal, Pool, PoolRecord, PoolsRequest,
+            },
+            poolmanager::v1beta1::SpotPriceRequest,
+        },
     },
-    osmosis::poolmanager::v1beta1::SpotPriceRequest,
 };
 use osmosis_test_tube::{
     cosmrs::proto::traits::Message,
@@ -119,6 +124,7 @@ pub fn fixture_dex_router(
             Uint128::zero(),
             Uint128::zero(),
             performance_fee,
+            true
         )
 }
 
@@ -257,7 +263,7 @@ pub fn init_test_contract(
 }
 
 #[allow(clippy::too_many_arguments)]
-fn init_test_contract_with_dex_router_and_swap_pools(
+pub fn init_test_contract_with_dex_router_and_swap_pools(
     filename_cl: &str,
     filename_dex: &str,
     admin_balance: &[Coin],
@@ -268,6 +274,7 @@ fn init_test_contract_with_dex_router_and_swap_pools(
     token_min_amount0: Uint128,
     token_min_amount1: Uint128,
     performance_fee: u64,
+    use_pool_coins: bool,
 ) -> (
     OsmosisTestApp,
     Addr,
@@ -361,9 +368,11 @@ fn init_test_contract_with_dex_router_and_swap_pools(
     // Create Balancer pools with previous vec of vec of coins
     // TODO: We could be using a mixed set of CL and gAMM pools here
     let mut lp_pools = vec![];
-    for pool_coins in &pools_coins {
-        let lp_pool = gm.create_basic_pool(pool_coins, &admin).unwrap();
-        lp_pools.push(lp_pool.data.pool_id);
+    if use_pool_coins {
+        for pool_coins in &pools_coins {
+            let lp_pool = gm.create_basic_pool(pool_coins, &admin).unwrap();
+            lp_pools.push(lp_pool.data.pool_id);
+        }
     }
 
     // Here we have 4 pools in pools_ids where the index 0 is the cl_pool id
@@ -399,7 +408,7 @@ fn init_test_contract_with_dex_router_and_swap_pools(
 
     let deposit_ratio_base = calculate_deposit_ratio(
         spot_price.spot_price,
-        tokens_provided,
+        tokens_provided.clone(),
         create_position.data.amount0,
         create_position.data.amount1,
         DENOM_BASE.to_string(),
@@ -422,13 +431,23 @@ fn init_test_contract_with_dex_router_and_swap_pools(
         .unwrap();
 
     // Here we pass only the 3x swap LP pools, not the Vault CL pool id 1
-    set_dex_router_paths(
-        &app,
-        &contract_dex_router.data.address,
-        &lp_pools,
-        &pools_coins,
-        &admin,
-    );
+    if use_pool_coins {
+        set_dex_router_paths(
+            &app,
+            &contract_dex_router.data.address,
+            &lp_pools,
+            &pools_coins,
+            &admin,
+        );
+    } else {
+        set_dex_router_paths(
+            &app,
+            &contract_dex_router.data.address,
+            &[vault_pool.id],
+            &[try_proto_to_cosmwasm_coins(tokens_provided).unwrap()],
+            &admin,
+        );
+    }
 
     // Instantiate vault
     let contract_cl = wasm