Migrate to PoolManager for Pool and SpotPrice. Read StableSwap pool.

Cargo.toml

@@ -62,7 +62,7 @@ neutron-sdk       = "0.6.0"
 # dev-dependencies
 cw-multi-test     = "0.16.5"
 cw-it             = "0.1.0"
-osmosis-test-tube = "16.0.0"
+osmosis-test-tube = "=16.1.1"
 test-case         = "3.0.0"
 proptest          = "1.1.0"
 

contracts/oracle/osmosis/src/helpers.rs

@@ -1,36 +1,71 @@
 use mars_oracle_base::{ContractError, ContractResult};
-use mars_osmosis::helpers::{has_denom, Pool};
+use mars_osmosis::{
+    helpers::{CommonPoolData, Pool},
+    BalancerPool,
+};
 
 use crate::DowntimeDetector;
 
 /// 48 hours in seconds
 const TWO_DAYS_IN_SECONDS: u64 = 172800u64;
 
-/// Assert the Osmosis pool indicated by `pool_id` is of XYK type and assets are OSMO and `denom`
+/// Assert the Osmosis pool indicated by `pool_id` is of Balancer XYK or StableSwap and assets are OSMO and `denom`
 pub fn assert_osmosis_pool_assets(
     pool: &Pool,
     denom: &str,
     base_denom: &str,
 ) -> ContractResult<()> {
-    assert_osmosis_xyk_pool(pool)?;
+    match pool {
+        Pool::Balancer(balancer_pool) => {
+            assert_osmosis_xyk_pool(balancer_pool)?;
+        }
+        Pool::StableSwap(_) => {}
+    };
+
+    assert_osmosis_pool_contains_two_assets(pool, denom, base_denom)?;
+
+    Ok(())
+}
+
+/// Assert the Osmosis pool indicated by `pool_id` is Balancer XYK type
+pub fn assert_osmosis_xyk_lp_pool(pool: &Pool) -> ContractResult<()> {
+    match pool {
+        Pool::Balancer(balancer_pool) => assert_osmosis_xyk_pool(balancer_pool)?,
+        Pool::StableSwap(stable_swap_pool) => {
+            return Err(ContractError::InvalidPriceSource {
+                reason: format!("StableSwap pool not supported. Pool id {}", stable_swap_pool.id),
+            });
+        }
+    };
+
+    Ok(())
+}
+
+fn assert_osmosis_pool_contains_two_assets(
+    pool: &Pool,
+    denom: &str,
+    base_denom: &str,
+) -> ContractResult<()> {
+    let pool_id = pool.get_pool_id();
+    let pool_denoms = pool.get_pool_denoms();
 
-    if !has_denom(base_denom, &pool.pool_assets) {
+    if !pool_denoms.contains(&base_denom.to_string()) {
         return Err(ContractError::InvalidPriceSource {
-            reason: format!("pool {} does not contain the base denom {}", pool.id, base_denom),
+            reason: format!("pool {} does not contain the base denom {}", pool_id, base_denom),
         });
     }
 
-    if !has_denom(denom, &pool.pool_assets) {
+    if !pool_denoms.contains(&denom.to_string()) {
         return Err(ContractError::InvalidPriceSource {
-            reason: format!("pool {} does not contain {}", pool.id, denom),
+            reason: format!("pool {} does not contain {}", pool_id, denom),
         });
     }
 
     Ok(())
 }
 
 /// Assert the Osmosis pool indicated by `pool_id` is of XYK type
-pub fn assert_osmosis_xyk_pool(pool: &Pool) -> ContractResult<()> {
+pub fn assert_osmosis_xyk_pool(pool: &BalancerPool) -> ContractResult<()> {
     if pool.pool_assets.len() != 2 {
         return Err(ContractError::InvalidPriceSource {
             reason: format!(

contracts/oracle/osmosis/src/price_source.rs

@@ -3,7 +3,8 @@ use std::{cmp::min, fmt};
 use cosmwasm_std::{Addr, Decimal, Decimal256, Deps, Empty, Env, Isqrt, Uint128, Uint256};
 use cw_storage_plus::Map;
 use mars_oracle_base::{
-    ContractError::InvalidPrice, ContractResult, PriceSourceChecked, PriceSourceUnchecked,
+    ContractError::{self, InvalidPrice},
+    ContractResult, PriceSourceChecked, PriceSourceUnchecked,
 };
 use mars_osmosis::helpers::{
     query_arithmetic_twap_price, query_geometric_twap_price, query_pool, query_spot_price,
@@ -369,7 +370,7 @@ impl PriceSourceUnchecked<OsmosisPriceSourceChecked, Empty> for OsmosisPriceSour
                 pool_id,
             } => {
                 let pool = query_pool(&deps.querier, *pool_id)?;
-                helpers::assert_osmosis_xyk_pool(&pool)?;
+                helpers::assert_osmosis_xyk_lp_pool(&pool)?;
                 Ok(OsmosisPriceSourceChecked::XykLiquidityToken {
                     pool_id: *pool_id,
                 })
@@ -593,6 +594,14 @@ impl OsmosisPriceSourceChecked {
     ) -> ContractResult<Decimal> {
         // XYK pool asserted during price source creation
         let pool = query_pool(&deps.querier, pool_id)?;
+        let pool = match pool {
+            Pool::Balancer(pool) => pool,
+            Pool::StableSwap(pool) => {
+                return Err(ContractError::InvalidPrice {
+                    reason: format!("StableSwap pool not supported. Pool id {}", pool.id),
+                })
+            }
+        };
 
         let coin0 = Pool::unwrap_coin(&pool.pool_assets[0].token)?;
         let coin1 = Pool::unwrap_coin(&pool.pool_assets[1].token)?;

contracts/oracle/osmosis/tests/helpers.rs

@@ -9,10 +9,10 @@ use cosmwasm_std::{
 };
 use mars_oracle_base::ContractError;
 use mars_oracle_osmosis::{contract::entry, msg::ExecuteMsg, OsmosisPriceSourceUnchecked};
-use mars_osmosis::helpers::{Pool, QueryPoolResponse};
+use mars_osmosis::{BalancerPool, StableSwapPool};
 use mars_red_bank_types::oracle::msg::{InstantiateMsg, QueryMsg};
 use mars_testing::{mock_info, MarsMockQuerier};
-use osmosis_std::types::osmosis::gamm::v1beta1::PoolAsset;
+use osmosis_std::types::osmosis::{gamm::v1beta1::PoolAsset, poolmanager::v1beta1::PoolResponse};
 use pyth_sdk_cw::PriceIdentifier;
 
 pub fn setup_test_with_pools() -> OwnedDeps<MockStorage, MockApi, MarsMockQuerier> {
@@ -22,25 +22,40 @@ pub fn setup_test_with_pools() -> OwnedDeps<MockStorage, MockApi, MarsMockQuerie
     let assets = vec![coin(42069, "uatom"), coin(69420, "uosmo")];
     deps.querier.set_query_pool_response(
         1,
-        prepare_query_pool_response(1, &assets, &[5000u64, 5000u64], &coin(10000, "gamm/pool/1")),
+        prepare_query_balancer_pool_response(
+            1,
+            &assets,
+            &[5000u64, 5000u64],
+            &coin(10000, "gamm/pool/1"),
+        ),
     );
 
     let assets = vec![coin(12345, "uusdc"), coin(23456, "uatom")];
     deps.querier.set_query_pool_response(
         64,
-        prepare_query_pool_response(64, &assets, &[5000u64, 5000u64], &coin(10000, "gamm/pool/64")),
+        prepare_query_balancer_pool_response(
+            64,
+            &assets,
+            &[5000u64, 5000u64],
+            &coin(10000, "gamm/pool/64"),
+        ),
     );
 
     let assets = vec![coin(12345, "uosmo"), coin(88888, "umars")];
     deps.querier.set_query_pool_response(
         89,
-        prepare_query_pool_response(89, &assets, &[5000u64, 5000u64], &coin(10000, "gamm/pool/89")),
+        prepare_query_balancer_pool_response(
+            89,
+            &assets,
+            &[5000u64, 5000u64],
+            &coin(10000, "gamm/pool/89"),
+        ),
     );
 
     let assets = vec![coin(12345, "ustatom"), coin(88888, "uatom")];
     deps.querier.set_query_pool_response(
         803,
-        prepare_query_pool_response(
+        prepare_query_balancer_pool_response(
             803,
             &assets,
             &[5000u64, 5000u64],
@@ -51,7 +66,7 @@ pub fn setup_test_with_pools() -> OwnedDeps<MockStorage, MockApi, MarsMockQuerie
     let assets = vec![coin(100000, "uusdc"), coin(100000, "uusdt"), coin(100000, "udai")];
     deps.querier.set_query_pool_response(
         3333,
-        prepare_query_pool_response(
+        prepare_query_balancer_pool_response(
             3333,
             &assets,
             &[5000u64, 5000u64, 5000u64],
@@ -63,14 +78,19 @@ pub fn setup_test_with_pools() -> OwnedDeps<MockStorage, MockApi, MarsMockQuerie
     let assets = vec![coin(100000, "uion"), coin(100000, "uosmo")];
     deps.querier.set_query_pool_response(
         4444,
-        prepare_query_pool_response(
+        prepare_query_balancer_pool_response(
             4444,
             &assets,
             &[5000u64, 5005u64],
             &coin(10000, "gamm/pool/4444"),
         ),
     );
 
+    // Set StableSwap pool
+    let assets = vec![coin(42069, "uatom"), coin(69420, "uosmo")];
+    deps.querier
+        .set_query_pool_response(5555, prepare_query_stable_swap_pool_response(5555, &assets));
+
     deps
 }
 
@@ -113,15 +133,15 @@ pub fn setup_test() -> OwnedDeps<MockStorage, MockApi, MarsMockQuerier> {
     deps
 }
 
-pub fn prepare_query_pool_response(
+pub fn prepare_query_balancer_pool_response(
     pool_id: u64,
     assets: &[Coin],
     weights: &[u64],
     shares: &Coin,
-) -> QueryPoolResponse {
-    let pool = Pool {
+) -> PoolResponse {
+    let pool = BalancerPool {
         address: "address".to_string(),
-        id: pool_id.to_string(),
+        id: pool_id,
         pool_params: None,
         future_pool_governor: "future_pool_governor".to_string(),
         total_shares: Some(osmosis_std::types::cosmos::base::v1beta1::Coin {
@@ -131,8 +151,8 @@ pub fn prepare_query_pool_response(
         pool_assets: prepare_pool_assets(assets, weights),
         total_weight: "".to_string(),
     };
-    QueryPoolResponse {
-        pool,
+    PoolResponse {
+        pool: Some(pool.to_any()),
     }
 }
 
@@ -155,6 +175,33 @@ fn prepare_pool_assets(coins: &[Coin], weights: &[u64]) -> Vec<PoolAsset> {
         .collect()
 }
 
+pub fn prepare_query_stable_swap_pool_response(pool_id: u64, assets: &[Coin]) -> PoolResponse {
+    let pool_liquidity: Vec<_> = assets
+        .iter()
+        .map(|coin| osmosis_std::types::cosmos::base::v1beta1::Coin {
+            denom: coin.denom.clone(),
+            amount: coin.amount.to_string(),
+        })
+        .collect();
+
+    let pool = StableSwapPool {
+        address: "osmo15v4mn84s9flhzpstkf9ql2mu0rnxh42pm8zhq47kh2fzs5zlwjsqaterkr".to_string(),
+        id: pool_id,
+        pool_params: None,
+        future_pool_governor: "".to_string(),
+        total_shares: Some(osmosis_std::types::cosmos::base::v1beta1::Coin {
+            denom: format!("gamm/pool/{pool_id}"),
+            amount: 4497913440357232330148u128.to_string(),
+        }),
+        pool_liquidity,
+        scaling_factors: vec![100000u64, 113890u64],
+        scaling_factor_controller: "osmo1k8c2m5cn322akk5wy8lpt87dd2f4yh9afcd7af".to_string(),
+    };
+    PoolResponse {
+        pool: Some(pool.to_any()),
+    }
+}
+
 pub fn set_pyth_price_source(deps: DepsMut, denom: &str, price_id: PriceIdentifier) {
     set_price_source(
         deps,

contracts/oracle/osmosis/tests/test_query_price.rs

@@ -18,7 +18,7 @@ use osmosis_std::types::osmosis::{
 };
 use pyth_sdk_cw::{Price, PriceFeed, PriceFeedResponse, PriceIdentifier};
 
-use crate::helpers::prepare_query_pool_response;
+use crate::helpers::prepare_query_balancer_pool_response;
 
 mod helpers;
 
@@ -804,7 +804,7 @@ fn querying_xyk_lp_price() {
     let assets = vec![coin(1, "uatom"), coin(1, "uosmo")];
     deps.querier.set_query_pool_response(
         10001,
-        prepare_query_pool_response(
+        prepare_query_balancer_pool_response(
             10001,
             &assets,
             &[5000u64, 5000u64],
@@ -815,7 +815,7 @@ fn querying_xyk_lp_price() {
     let assets = vec![coin(1, "umars"), coin(1, "uosmo")];
     deps.querier.set_query_pool_response(
         10002,
-        prepare_query_pool_response(
+        prepare_query_balancer_pool_response(
             10002,
             &assets,
             &[5000u64, 5000u64],
@@ -826,7 +826,7 @@ fn querying_xyk_lp_price() {
     let assets = vec![coin(10000, "uatom"), coin(885000, "umars")];
     deps.querier.set_query_pool_response(
         10003,
-        prepare_query_pool_response(
+        prepare_query_balancer_pool_response(
             10003,
             &assets,
             &[5000u64, 5000u64],
@@ -889,7 +889,7 @@ fn querying_xyk_lp_price() {
     let assets = vec![coin(6389, "uatom"), coin(1385000, "umars")];
     deps.querier.set_query_pool_response(
         10003,
-        prepare_query_pool_response(
+        prepare_query_balancer_pool_response(
             10003,
             &assets,
             &[5000u64, 5000u64],

contracts/oracle/osmosis/tests/test_set_price_source.rs

@@ -926,6 +926,15 @@ fn setting_price_source_xyk_lp() {
         }
     );
 
+    // attempting to use StableSwap pool
+    let err = set_price_source_xyk_lp("atom_mars_lp", 5555).unwrap_err();
+    assert_eq!(
+        err,
+        ContractError::InvalidPriceSource {
+            reason: "StableSwap pool not supported. Pool id 5555".to_string()
+        }
+    );
+
     // properly set xyk lp price source
     let res = set_price_source_xyk_lp("uosmo_umars_lp", 89).unwrap();
     assert_eq!(res.messages.len(), 0);