osmo_bindings_test/
multitest.rs

1use anyhow::{bail, Result as AnyResult};
2use itertools::Itertools;
3use schemars::JsonSchema;
4use serde::de::DeserializeOwned;
5use serde::{Deserialize, Serialize};
6use std::cmp::max;
7use std::fmt::Debug;
8use std::iter;
9use std::ops::{Deref, DerefMut};
10use thiserror::Error;
11
12use cosmwasm_std::testing::{MockApi, MockStorage};
13use cosmwasm_std::{
14    coins, to_binary, Addr, Api, BankMsg, Binary, BlockInfo, Coin, CustomQuery, Decimal, Empty,
15    Fraction, Isqrt, Querier, QuerierResult, StdError, StdResult, Storage, Uint128,
16};
17use cw_multi_test::{
18    App, AppResponse, BankKeeper, BankSudo, BasicAppBuilder, CosmosRouter, Module, WasmKeeper,
19};
20use cw_storage_plus::Map;
21
22use crate::error::ContractError;
23use osmo_bindings::{
24    FullDenomResponse, OsmosisMsg, OsmosisQuery, PoolStateResponse, SpotPriceResponse, Step, Swap,
25    SwapAmount, SwapAmountWithLimit, SwapResponse,
26};
27
28pub const POOLS: Map<u64, Pool> = Map::new("pools");
29
30#[derive(Serialize, Deserialize, Clone, PartialEq, JsonSchema, Debug)]
31pub struct Pool {
32    pub assets: Vec<Coin>,
33    pub shares: Uint128,
34    pub fee: Decimal,
35}
36
37impl Pool {
38    // make an equal-weighted uniswap-like pool with 0.3% fees
39    pub fn new(a: Coin, b: Coin) -> Self {
40        let shares = (a.amount * b.amount).isqrt();
41        Pool {
42            assets: vec![a, b],
43            shares,
44            fee: Decimal::permille(3),
45        }
46    }
47
48    pub fn has_denom(&self, denom: &str) -> bool {
49        self.assets.iter().any(|c| c.denom == denom)
50    }
51
52    pub fn get_amount(&self, denom: &str) -> Option<Uint128> {
53        self.assets
54            .iter()
55            .find(|c| c.denom == denom)
56            .map(|c| c.amount)
57    }
58
59    pub fn set_amount(&mut self, denom: &str, amount: Uint128) -> Result<(), OsmosisError> {
60        let pos = self
61            .assets
62            .iter()
63            .position(|c| c.denom == denom)
64            .ok_or(OsmosisError::AssetNotInPool)?;
65        self.assets[pos].amount = amount;
66        Ok(())
67    }
68
69    pub fn spot_price(
70        &self,
71        denom_in: &str,
72        denom_out: &str,
73        with_swap_fee: bool,
74    ) -> Result<Decimal, OsmosisError> {
75        // ensure they have both assets
76        let (bal_in, bal_out) = match (self.get_amount(denom_in), self.get_amount(denom_out)) {
77            (Some(a), Some(b)) => (a, b),
78            _ => return Err(OsmosisError::AssetNotInPool),
79        };
80        let mult = if with_swap_fee {
81            Decimal::one() - self.fee
82        } else {
83            Decimal::one()
84        };
85        let price = Decimal::from_ratio(bal_out * mult, bal_in);
86        Ok(price)
87    }
88
89    pub fn swap(
90        &mut self,
91        denom_in: &str,
92        denom_out: &str,
93        amount: SwapAmount,
94    ) -> Result<SwapAmount, OsmosisError> {
95        // ensure they have both assets
96        let (bal_in, bal_out) = match (self.get_amount(denom_in), self.get_amount(denom_out)) {
97            (Some(a), Some(b)) => (a, b),
98            _ => return Err(OsmosisError::AssetNotInPool),
99        };
100        // do calculations (in * out = k) equation
101        let (final_in, final_out, payout) = match amount {
102            SwapAmount::In(input) => {
103                let input_minus_fee = input * (Decimal::one() - self.fee);
104                let final_out = bal_in * bal_out / (bal_in + input_minus_fee);
105                let payout = SwapAmount::Out(bal_out - final_out);
106                let final_in = bal_in + input;
107                (final_in, final_out, payout)
108            }
109            SwapAmount::Out(output) => {
110                let in_without_fee = bal_in * bal_out / (bal_out - output);
111                // add one to handle rounding (final_in - old_in) / (1 - fee)
112                let mult = Decimal::one() - self.fee;
113                // Use this as Uint128 / Decimal is not implemented in cosmwasm_std
114                let pay_incl_fee = (in_without_fee - bal_in) * mult.denominator()
115                    / mult.numerator()
116                    + Uint128::new(1);
117
118                let payin = SwapAmount::In(pay_incl_fee);
119                let final_in = bal_in + pay_incl_fee;
120                let final_out = bal_out - output;
121                (final_in, final_out, payin)
122            }
123        };
124        // update internal balance
125        self.set_amount(denom_in, final_in)?;
126        self.set_amount(denom_out, final_out)?;
127        Ok(payout)
128    }
129
130    pub fn swap_with_limit(
131        &mut self,
132        denom_in: &str,
133        denom_out: &str,
134        amount: SwapAmountWithLimit,
135    ) -> Result<SwapAmount, OsmosisError> {
136        match amount {
137            SwapAmountWithLimit::ExactIn { input, min_output } => {
138                let payout = self.swap(denom_in, denom_out, SwapAmount::In(input))?;
139                if payout.as_out() < min_output {
140                    Err(OsmosisError::PriceTooLow)
141                } else {
142                    Ok(payout)
143                }
144            }
145            SwapAmountWithLimit::ExactOut { output, max_input } => {
146                let payin = self.swap(denom_in, denom_out, SwapAmount::Out(output))?;
147                if payin.as_in() > max_input {
148                    Err(OsmosisError::PriceTooLow)
149                } else {
150                    Ok(payin)
151                }
152            }
153        }
154    }
155
156    pub fn gamm_denom(&self, pool_id: u64) -> String {
157        // see https://github.com/osmosis-labs/osmosis/blob/e13cddc698a121dce2f8919b2a0f6a743f4082d6/x/gamm/types/key.go#L52-L54
158        format!("gamm/pool/{}", pool_id)
159    }
160
161    pub fn into_response(self, pool_id: u64) -> PoolStateResponse {
162        let denom = self.gamm_denom(pool_id);
163        PoolStateResponse {
164            assets: self.assets,
165            shares: Coin {
166                denom,
167                amount: self.shares,
168            },
169        }
170    }
171}
172
173pub struct OsmosisModule {}
174
175/// How many seconds per block
176/// (when we increment block.height, use this multiplier for block.time)
177pub const BLOCK_TIME: u64 = 5;
178
179impl OsmosisModule {
180    fn build_denom(&self, contract: &Addr, sub_denom: &str) -> Result<String, ContractError> {
181        // Minimum validation checks on the full denom.
182        // https://github.com/cosmos/cosmos-sdk/blob/2646b474c7beb0c93d4fafd395ef345f41afc251/types/coin.go#L706-L711
183        // https://github.com/cosmos/cosmos-sdk/blob/2646b474c7beb0c93d4fafd395ef345f41afc251/types/coin.go#L677
184        let full_denom = format!("factory/{}/{}", contract, sub_denom);
185        if full_denom.len() < 3 || full_denom.len() > 128 || contract.as_str().contains('/') {
186            return Err(ContractError::InvalidFullDenom { full_denom });
187        }
188        Ok(full_denom)
189    }
190
191    /// Used to mock out the response for TgradeQuery::ValidatorVotes
192    pub fn set_pool(&self, storage: &mut dyn Storage, pool_id: u64, pool: &Pool) -> StdResult<()> {
193        POOLS.save(storage, pool_id, pool)
194    }
195}
196
197fn complex_swap(
198    storage: &dyn Storage,
199    first: Swap,
200    route: Vec<Step>,
201    amount: SwapAmount,
202) -> AnyResult<(SwapAmount, Vec<(u64, Pool)>)> {
203    // all the `Swap`s we need to execute in order
204    let swaps: Vec<_> = {
205        let frst = iter::once(first.clone());
206        let rest = iter::once((first.pool_id, first.denom_out))
207            .chain(route.into_iter().map(|step| (step.pool_id, step.denom_out)))
208            .tuple_windows()
209            .map(|((_, denom_in), (pool_id, denom_out))| Swap {
210                pool_id,
211                denom_in,
212                denom_out,
213            });
214        frst.chain(rest).collect()
215    };
216
217    let mut updated_pools = vec![];
218
219    match amount {
220        SwapAmount::In(mut input) => {
221            for swap in &swaps {
222                let mut pool = POOLS.load(storage, swap.pool_id)?;
223                let payout = pool.swap(&swap.denom_in, &swap.denom_out, SwapAmount::In(input))?;
224                updated_pools.push((swap.pool_id, pool));
225
226                input = payout.as_out();
227            }
228
229            Ok((SwapAmount::Out(input), updated_pools))
230        }
231        SwapAmount::Out(mut output) => {
232            for swap in swaps.iter().rev() {
233                let mut pool = POOLS.load(storage, swap.pool_id)?;
234                let payout = pool.swap(&swap.denom_in, &swap.denom_out, SwapAmount::Out(output))?;
235                updated_pools.push((swap.pool_id, pool));
236
237                output = payout.as_in();
238            }
239
240            Ok((SwapAmount::In(output), updated_pools))
241        }
242    }
243}
244
245impl Module for OsmosisModule {
246    type ExecT = OsmosisMsg;
247    type QueryT = OsmosisQuery;
248    type SudoT = Empty;
249
250    fn execute<ExecC, QueryC>(
251        &self,
252        api: &dyn Api,
253        storage: &mut dyn Storage,
254        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
255        block: &BlockInfo,
256        sender: Addr,
257        msg: OsmosisMsg,
258    ) -> AnyResult<AppResponse>
259    where
260        ExecC: Debug + Clone + PartialEq + JsonSchema + DeserializeOwned + 'static,
261        QueryC: CustomQuery + DeserializeOwned + 'static,
262    {
263        match msg {
264            OsmosisMsg::MintTokens {
265                sub_denom,
266                amount,
267                recipient,
268            } => {
269                let denom = self.build_denom(&sender, &sub_denom)?;
270                let mint = BankSudo::Mint {
271                    to_address: recipient,
272                    amount: coins(amount.u128(), &denom),
273                };
274                router.sudo(api, storage, block, mint.into())?;
275
276                let data = Some(to_binary(&FullDenomResponse { denom })?);
277                Ok(AppResponse {
278                    data,
279                    events: vec![],
280                })
281            }
282            OsmosisMsg::Swap {
283                first,
284                route,
285                amount,
286            } => {
287                let denom_in = first.denom_in.clone();
288                let denom_out = route
289                    .iter()
290                    .last()
291                    .map(|step| step.denom_out.clone())
292                    .unwrap_or_else(|| first.denom_out.clone());
293
294                let (swap_result, updated_pools) =
295                    complex_swap(storage, first, route, amount.clone().discard_limit())?;
296
297                match amount {
298                    SwapAmountWithLimit::ExactIn { min_output, .. } => {
299                        if swap_result.as_out() < min_output {
300                            return Err(OsmosisError::PriceTooLow.into());
301                        }
302                    }
303                    SwapAmountWithLimit::ExactOut { max_input, .. } => {
304                        if swap_result.as_in() > max_input {
305                            return Err(OsmosisError::PriceTooLow.into());
306                        }
307                    }
308                }
309
310                for (pool_id, pool) in updated_pools {
311                    POOLS.save(storage, pool_id, &pool)?;
312                }
313
314                let (pay_in, get_out) = match amount {
315                    SwapAmountWithLimit::ExactIn { input, .. } => (input, swap_result.as_out()),
316                    SwapAmountWithLimit::ExactOut { output, .. } => (swap_result.as_in(), output),
317                };
318
319                // Note: to make testing easier, we just mint and burn - no balance for AMM
320                // burn pay_in tokens from sender
321                let burn = BankMsg::Burn {
322                    amount: coins(pay_in.u128(), &denom_in),
323                };
324                router.execute(api, storage, block, sender.clone(), burn.into())?;
325
326                // mint get_out tokens to sender
327                let mint = BankSudo::Mint {
328                    to_address: sender.to_string(),
329                    amount: coins(get_out.u128(), denom_out),
330                };
331                router.sudo(api, storage, block, mint.into())?;
332
333                let output = match amount {
334                    SwapAmountWithLimit::ExactIn { .. } => SwapAmount::Out(get_out),
335                    SwapAmountWithLimit::ExactOut { .. } => SwapAmount::In(pay_in),
336                };
337                let data = Some(to_binary(&SwapResponse { amount: output })?);
338                Ok(AppResponse {
339                    data,
340                    events: vec![],
341                })
342            }
343        }
344    }
345
346    fn sudo<ExecC, QueryC>(
347        &self,
348        _api: &dyn Api,
349        _storage: &mut dyn Storage,
350        _router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
351        _block: &BlockInfo,
352        _msg: Self::SudoT,
353    ) -> AnyResult<AppResponse>
354    where
355        ExecC: Debug + Clone + PartialEq + JsonSchema + DeserializeOwned + 'static,
356        QueryC: CustomQuery + DeserializeOwned + 'static,
357    {
358        bail!("sudo not implemented for OsmosisModule")
359    }
360
361    fn query(
362        &self,
363        api: &dyn Api,
364        storage: &dyn Storage,
365        _querier: &dyn Querier,
366        _block: &BlockInfo,
367        request: OsmosisQuery,
368    ) -> anyhow::Result<Binary> {
369        match request {
370            OsmosisQuery::FullDenom {
371                contract,
372                sub_denom,
373            } => {
374                let contract = api.addr_validate(&contract)?;
375                let denom = self.build_denom(&contract, &sub_denom)?;
376                let res = FullDenomResponse { denom };
377                Ok(to_binary(&res)?)
378            }
379            OsmosisQuery::PoolState { id } => {
380                let pool = POOLS.load(storage, id)?;
381                let res = pool.into_response(id);
382                Ok(to_binary(&res)?)
383            }
384            OsmosisQuery::SpotPrice {
385                swap,
386                with_swap_fee,
387            } => {
388                let pool = POOLS.load(storage, swap.pool_id)?;
389                let price = pool.spot_price(&swap.denom_in, &swap.denom_out, with_swap_fee)?;
390                Ok(to_binary(&SpotPriceResponse { price })?)
391            }
392            OsmosisQuery::EstimateSwap {
393                sender: _sender,
394                first,
395                route,
396                amount,
397            } => {
398                let (amount, _) = complex_swap(storage, first, route, amount)?;
399
400                Ok(to_binary(&SwapResponse { amount })?)
401            }
402        }
403    }
404}
405
406#[derive(Error, Debug, PartialEq)]
407pub enum OsmosisError {
408    #[error("{0}")]
409    Std(#[from] StdError),
410
411    #[error("Asset not in pool")]
412    AssetNotInPool,
413
414    #[error("Price under minimum requested, aborting swap")]
415    PriceTooLow,
416
417    /// Remove this to let the compiler find all TODOs
418    #[error("Not yet implemented (TODO)")]
419    Unimplemented,
420}
421
422pub type OsmosisAppWrapped =
423    App<BankKeeper, MockApi, MockStorage, OsmosisModule, WasmKeeper<OsmosisMsg, OsmosisQuery>>;
424
425pub struct OsmosisApp(OsmosisAppWrapped);
426
427impl Deref for OsmosisApp {
428    type Target = OsmosisAppWrapped;
429
430    fn deref(&self) -> &Self::Target {
431        &self.0
432    }
433}
434
435impl DerefMut for OsmosisApp {
436    fn deref_mut(&mut self) -> &mut Self::Target {
437        &mut self.0
438    }
439}
440
441impl Querier for OsmosisApp {
442    fn raw_query(&self, bin_request: &[u8]) -> QuerierResult {
443        self.0.raw_query(bin_request)
444    }
445}
446
447impl Default for OsmosisApp {
448    fn default() -> Self {
449        Self::new()
450    }
451}
452
453impl OsmosisApp {
454    pub fn new() -> Self {
455        Self(
456            BasicAppBuilder::<OsmosisMsg, OsmosisQuery>::new_custom()
457                .with_custom(OsmosisModule {})
458                .build(|_router, _, _storage| {
459                    // router.custom.set_owner(storage, &owner).unwrap();
460                }),
461        )
462    }
463
464    pub fn block_info(&self) -> BlockInfo {
465        self.0.block_info()
466    }
467
468    /// This advances BlockInfo by given number of blocks.
469    /// It does not do any callbacks, but keeps the ratio of seconds/blokc
470    pub fn advance_blocks(&mut self, blocks: u64) {
471        self.update_block(|block| {
472            block.time = block.time.plus_seconds(BLOCK_TIME * blocks);
473            block.height += blocks;
474        });
475    }
476
477    /// This advances BlockInfo by given number of seconds.
478    /// It does not do any callbacks, but keeps the ratio of seconds/blokc
479    pub fn advance_seconds(&mut self, seconds: u64) {
480        self.update_block(|block| {
481            block.time = block.time.plus_seconds(seconds);
482            block.height += max(1, seconds / BLOCK_TIME);
483        });
484    }
485
486    /// Simple iterator when you don't care too much about the details and just want to
487    /// simulate forward motion.
488    pub fn next_block(&mut self) {
489        self.advance_blocks(1)
490    }
491}
492
493#[cfg(test)]
494mod tests {
495    use super::*;
496    use cosmwasm_std::testing::MOCK_CONTRACT_ADDR;
497    use cosmwasm_std::{coin, from_slice, Uint128};
498    use cw_multi_test::Executor;
499    use osmo_bindings::{Step, Swap};
500
501    #[test]
502    fn mint_token() {
503        let contract = Addr::unchecked("govner");
504        let rcpt = Addr::unchecked("townies");
505        let sub_denom = "fundz";
506
507        let mut app = OsmosisApp::new();
508
509        // no tokens
510        let start = app.wrap().query_all_balances(rcpt.as_str()).unwrap();
511        assert_eq!(start, vec![]);
512
513        // let's find the mapping
514        let FullDenomResponse { denom } = app
515            .wrap()
516            .query(
517                &OsmosisQuery::FullDenom {
518                    contract: contract.to_string(),
519                    sub_denom: sub_denom.to_string(),
520                }
521                .into(),
522            )
523            .unwrap();
524        assert_ne!(denom, sub_denom);
525        assert!(denom.len() > 10);
526
527        // prepare to mint
528        let amount = Uint128::new(1234567);
529        let msg = OsmosisMsg::MintTokens {
530            sub_denom: sub_denom.to_string(),
531            amount,
532            recipient: rcpt.to_string(),
533        };
534
535        // simulate contract calling
536        app.execute(contract, msg.into()).unwrap();
537
538        // we got tokens!
539        let end = app.wrap().query_balance(rcpt.as_str(), &denom).unwrap();
540        let expected = Coin { denom, amount };
541        assert_eq!(end, expected);
542
543        // but no minting of unprefixed version
544        let empty = app.wrap().query_balance(rcpt.as_str(), sub_denom).unwrap();
545        assert_eq!(empty.amount, Uint128::zero());
546    }
547
548    #[test]
549    fn query_pool() {
550        let coin_a = coin(6_000_000u128, "osmo");
551        let coin_b = coin(1_500_000u128, "atom");
552        let pool_id = 43;
553        let pool = Pool::new(coin_a.clone(), coin_b.clone());
554
555        // set up with one pool
556        let mut app = OsmosisApp::new();
557        app.init_modules(|router, _, storage| {
558            router.custom.set_pool(storage, pool_id, &pool).unwrap();
559        });
560
561        // query the pool state
562        let query = OsmosisQuery::PoolState { id: pool_id }.into();
563        let state: PoolStateResponse = app.wrap().query(&query).unwrap();
564        let expected_shares = coin(3_000_000, "gamm/pool/43");
565        assert_eq!(state.shares, expected_shares);
566        assert_eq!(state.assets, vec![coin_a.clone(), coin_b.clone()]);
567
568        // check spot price both directions
569        let query = OsmosisQuery::spot_price(pool_id, &coin_a.denom, &coin_b.denom).into();
570        let SpotPriceResponse { price } = app.wrap().query(&query).unwrap();
571        assert_eq!(price, Decimal::percent(25));
572
573        // and atom -> osmo
574        let query = OsmosisQuery::spot_price(pool_id, &coin_b.denom, &coin_a.denom).into();
575        let SpotPriceResponse { price } = app.wrap().query(&query).unwrap();
576        assert_eq!(price, Decimal::percent(400));
577
578        // with fee
579        let query = OsmosisQuery::SpotPrice {
580            swap: Swap::new(pool_id, &coin_b.denom, &coin_a.denom),
581            with_swap_fee: true,
582        };
583        let SpotPriceResponse { price } = app.wrap().query(&query.into()).unwrap();
584        // 4.00 * (1- 0.3%) = 4 * 0.997 = 3.988
585        assert_eq!(price, Decimal::permille(3988));
586    }
587
588    #[test]
589    fn estimate_swap() {
590        let coin_a = coin(6_000_000u128, "osmo");
591        let coin_b = coin(1_500_000u128, "atom");
592        let pool_id = 43;
593        let pool = Pool::new(coin_a.clone(), coin_b.clone());
594
595        // set up with one pool
596        let mut app = OsmosisApp::new();
597        app.init_modules(|router, _, storage| {
598            router.custom.set_pool(storage, pool_id, &pool).unwrap();
599        });
600
601        // estimate the price (501505 * 0.997 = 500_000) after fees gone
602        let query = OsmosisQuery::estimate_swap(
603            MOCK_CONTRACT_ADDR,
604            pool_id,
605            &coin_b.denom,
606            &coin_a.denom,
607            SwapAmount::In(Uint128::new(501505)),
608        );
609        let SwapResponse { amount } = app.wrap().query(&query.into()).unwrap();
610        // 6M * 1.5M = 2M * 4.5M -> output = 1.5M
611        let expected = SwapAmount::Out(Uint128::new(1_500_000));
612        assert_eq!(amount, expected);
613
614        // now try the reverse query. we know what we need to pay to get 1.5M out
615        let query = OsmosisQuery::estimate_swap(
616            MOCK_CONTRACT_ADDR,
617            pool_id,
618            &coin_b.denom,
619            &coin_a.denom,
620            SwapAmount::Out(Uint128::new(1500000)),
621        );
622        let SwapResponse { amount } = app.wrap().query(&query.into()).unwrap();
623        let expected = SwapAmount::In(Uint128::new(501505));
624        assert_eq!(amount, expected);
625    }
626
627    #[test]
628    fn perform_swap() {
629        let coin_a = coin(6_000_000u128, "osmo");
630        let coin_b = coin(1_500_000u128, "atom");
631        let pool_id = 43;
632        let pool = Pool::new(coin_a.clone(), coin_b.clone());
633        let trader = Addr::unchecked("trader");
634
635        // set up with one pool
636        let mut app = OsmosisApp::new();
637        app.init_modules(|router, _, storage| {
638            router.custom.set_pool(storage, pool_id, &pool).unwrap();
639            router
640                .bank
641                .init_balance(storage, &trader, coins(800_000, &coin_b.denom))
642                .unwrap()
643        });
644
645        // check balance before
646        let Coin { amount, .. } = app.wrap().query_balance(&trader, &coin_a.denom).unwrap();
647        assert_eq!(amount, Uint128::new(0));
648        let Coin { amount, .. } = app.wrap().query_balance(&trader, &coin_b.denom).unwrap();
649        assert_eq!(amount, Uint128::new(800_000));
650
651        // this is too low a payment, will error
652        let msg = OsmosisMsg::simple_swap(
653            pool_id,
654            &coin_b.denom,
655            &coin_a.denom,
656            SwapAmountWithLimit::ExactOut {
657                output: Uint128::new(1_500_000),
658                max_input: Uint128::new(400_000),
659            },
660        );
661        let err = app.execute(trader.clone(), msg.into()).unwrap_err();
662        println!("{:?}", err);
663
664        // now a proper swap
665        let msg = OsmosisMsg::simple_swap(
666            pool_id,
667            &coin_b.denom,
668            &coin_a.denom,
669            SwapAmountWithLimit::ExactOut {
670                output: Uint128::new(1_500_000),
671                max_input: Uint128::new(600_000),
672            },
673        );
674        let res = app.execute(trader.clone(), msg.into()).unwrap();
675
676        // update balances (800_000 - 501_505 paid = 298_495)
677        let Coin { amount, .. } = app.wrap().query_balance(&trader, &coin_a.denom).unwrap();
678        assert_eq!(amount, Uint128::new(1_500_000));
679        let Coin { amount, .. } = app.wrap().query_balance(&trader, &coin_b.denom).unwrap();
680        assert_eq!(amount, Uint128::new(298_495));
681
682        // check the response contains proper value
683        let input: SwapResponse = from_slice(res.data.unwrap().as_slice()).unwrap();
684        assert_eq!(input.amount, SwapAmount::In(Uint128::new(501_505)));
685
686        // check pool state properly updated with fees
687        let query = OsmosisQuery::PoolState { id: pool_id }.into();
688        let state: PoolStateResponse = app.wrap().query(&query).unwrap();
689        let expected_assets = vec![
690            coin(4_500_000, &coin_a.denom),
691            coin(2_001_505, &coin_b.denom),
692        ];
693        assert_eq!(state.assets, expected_assets);
694    }
695
696    #[test]
697    fn swap_with_route_max_input_exceeded() {
698        let pool1 = Pool::new(coin(6_000_000, "osmo"), coin(3_000_000, "atom"));
699        let pool2 = Pool::new(coin(2_000_000, "atom"), coin(1_000_000, "btc"));
700        let trader = Addr::unchecked("trader");
701
702        let mut app = OsmosisApp::new();
703        app.init_modules(|router, _, storage| {
704            router.custom.set_pool(storage, 1, &pool1).unwrap();
705            router.custom.set_pool(storage, 2, &pool2).unwrap();
706            router
707                .bank
708                .init_balance(storage, &trader, coins(5000, "osmo"))
709                .unwrap()
710        });
711
712        let msg = OsmosisMsg::Swap {
713            first: Swap {
714                pool_id: 1,
715                denom_in: "osmo".to_string(),
716                denom_out: "atom".to_string(),
717            },
718            route: vec![Step {
719                pool_id: 2,
720                denom_out: "btc".to_string(),
721            }],
722            amount: SwapAmountWithLimit::ExactOut {
723                output: Uint128::new(1000),
724                max_input: Uint128::new(4000),
725            },
726        };
727        let err = app.execute(trader, msg.into()).unwrap_err();
728        assert_eq!(
729            err.downcast::<OsmosisError>().unwrap(),
730            OsmosisError::PriceTooLow
731        );
732    }
733
734    #[test]
735    fn swap_with_route_min_output_not_met() {
736        let pool1 = Pool::new(coin(6_000_000, "osmo"), coin(3_000_000, "atom"));
737        let pool2 = Pool::new(coin(2_000_000, "atom"), coin(1_000_000, "btc"));
738        let trader = Addr::unchecked("trader");
739
740        let mut app = OsmosisApp::new();
741        app.init_modules(|router, _, storage| {
742            router.custom.set_pool(storage, 1, &pool1).unwrap();
743            router.custom.set_pool(storage, 2, &pool2).unwrap();
744            router
745                .bank
746                .init_balance(storage, &trader, coins(5000, "osmo"))
747                .unwrap()
748        });
749
750        let msg = OsmosisMsg::Swap {
751            first: Swap {
752                pool_id: 1,
753                denom_in: "osmo".to_string(),
754                denom_out: "atom".to_string(),
755            },
756            route: vec![Step {
757                pool_id: 2,
758                denom_out: "btc".to_string(),
759            }],
760            amount: SwapAmountWithLimit::ExactIn {
761                input: Uint128::new(4000),
762                min_output: Uint128::new(1000),
763            },
764        };
765        let err = app.execute(trader, msg.into()).unwrap_err();
766        assert_eq!(
767            err.downcast::<OsmosisError>().unwrap(),
768            OsmosisError::PriceTooLow
769        );
770    }
771
772    #[test]
773    fn swap_with_route_wrong_denom() {
774        let pool1 = Pool::new(coin(6_000_000, "osmo"), coin(3_000_000, "atom"));
775        let pool2 = Pool::new(coin(2_000_000, "atom"), coin(1_000_000, "eth"));
776        let trader = Addr::unchecked("trader");
777
778        let mut app = OsmosisApp::new();
779        app.init_modules(|router, _, storage| {
780            router.custom.set_pool(storage, 1, &pool1).unwrap();
781            router.custom.set_pool(storage, 2, &pool2).unwrap();
782            router
783                .bank
784                .init_balance(storage, &trader, coins(5000, "osmo"))
785                .unwrap()
786        });
787
788        let msg = OsmosisMsg::Swap {
789            first: Swap {
790                pool_id: 1,
791                denom_in: "osmo".to_string(),
792                denom_out: "atom".to_string(),
793            },
794            route: vec![Step {
795                pool_id: 2,
796                denom_out: "btc".to_string(),
797            }],
798            amount: SwapAmountWithLimit::ExactOut {
799                output: Uint128::new(1000),
800                max_input: Uint128::new(4000),
801            },
802        };
803        let err = app.execute(trader, msg.into()).unwrap_err();
804        assert_eq!(
805            err.downcast::<OsmosisError>().unwrap(),
806            OsmosisError::AssetNotInPool
807        );
808    }
809
810    #[test]
811    fn perform_swap_with_route_exact_out() {
812        let pool1 = Pool::new(coin(6_000_000, "osmo"), coin(3_000_000, "atom"));
813        let pool2 = Pool::new(coin(2_000_000, "atom"), coin(1_000_000, "btc"));
814        let trader = Addr::unchecked("trader");
815
816        // set up pools
817        let mut app = OsmosisApp::new();
818        app.init_modules(|router, _, storage| {
819            router.custom.set_pool(storage, 1, &pool1).unwrap();
820            router.custom.set_pool(storage, 2, &pool2).unwrap();
821            router
822                .bank
823                .init_balance(storage, &trader, coins(5000, "osmo"))
824                .unwrap()
825        });
826
827        let msg = OsmosisMsg::Swap {
828            first: Swap {
829                pool_id: 1,
830                denom_in: "osmo".to_string(),
831                denom_out: "atom".to_string(),
832            },
833            route: vec![Step {
834                pool_id: 2,
835                denom_out: "btc".to_string(),
836            }],
837            amount: SwapAmountWithLimit::ExactOut {
838                output: Uint128::new(1000),
839                max_input: Uint128::new(5000),
840            },
841        };
842        let res = app.execute(trader.clone(), msg.into()).unwrap();
843
844        let Coin { amount, .. } = app.wrap().query_balance(&trader, "osmo").unwrap();
845        assert_eq!(amount, Uint128::new(5000 - 4033));
846        let Coin { amount, .. } = app.wrap().query_balance(&trader, "btc").unwrap();
847        assert_eq!(amount, Uint128::new(1000));
848
849        // check the response contains proper value
850        let input: SwapResponse = from_slice(res.data.unwrap().as_slice()).unwrap();
851        assert_eq!(input.amount, SwapAmount::In(Uint128::new(4033)));
852
853        // check pool state properly updated with fees
854        let query = OsmosisQuery::PoolState { id: 1 }.into();
855        let state: PoolStateResponse = app.wrap().query(&query).unwrap();
856        let expected_assets = vec![
857            coin(6_000_000 + 4033, "osmo"),
858            coin(3_000_000 - 2009, "atom"),
859        ];
860        assert_eq!(state.assets, expected_assets);
861
862        let query = OsmosisQuery::PoolState { id: 2 }.into();
863        let state: PoolStateResponse = app.wrap().query(&query).unwrap();
864        let expected_assets = vec![
865            coin(2_000_000 + 2009, "atom"),
866            coin(1_000_000 - 1000, "btc"),
867        ];
868        assert_eq!(state.assets, expected_assets);
869    }
870
871    #[test]
872    fn perform_swap_with_route_exact_in() {
873        let pool1 = Pool::new(coin(6_000_000, "osmo"), coin(3_000_000, "atom"));
874        let pool2 = Pool::new(coin(2_000_000, "atom"), coin(1_000_000, "btc"));
875        let trader = Addr::unchecked("trader");
876
877        // set up pools
878        let mut app = OsmosisApp::new();
879        app.init_modules(|router, _, storage| {
880            router.custom.set_pool(storage, 1, &pool1).unwrap();
881            router.custom.set_pool(storage, 2, &pool2).unwrap();
882            router
883                .bank
884                .init_balance(storage, &trader, coins(5000, "osmo"))
885                .unwrap()
886        });
887
888        // now a proper swap
889        let msg = OsmosisMsg::Swap {
890            first: Swap {
891                pool_id: 1,
892                denom_in: "osmo".to_string(),
893                denom_out: "atom".to_string(),
894            },
895            route: vec![Step {
896                pool_id: 2,
897                denom_out: "btc".to_string(),
898            }],
899            amount: SwapAmountWithLimit::ExactIn {
900                input: Uint128::new(4000),
901                min_output: Uint128::new(900),
902            },
903        };
904        let res = app.execute(trader.clone(), msg.into()).unwrap();
905
906        let Coin { amount, .. } = app.wrap().query_balance(&trader, "osmo").unwrap();
907        assert_eq!(amount, Uint128::new(5000 - 4000));
908        let Coin { amount, .. } = app.wrap().query_balance(&trader, "btc").unwrap();
909        assert_eq!(amount, Uint128::new(993));
910
911        // check the response contains proper value
912        let input: SwapResponse = from_slice(res.data.unwrap().as_slice()).unwrap();
913        assert_eq!(input.amount, SwapAmount::Out(Uint128::new(993)));
914
915        // check pool state properly updated with fees
916        let query = OsmosisQuery::PoolState { id: 1 }.into();
917        let state: PoolStateResponse = app.wrap().query(&query).unwrap();
918        let expected_assets = vec![
919            coin(6_000_000 + 4000, "osmo"),
920            coin(3_000_000 - 1993, "atom"),
921        ];
922        assert_eq!(state.assets, expected_assets);
923
924        let query = OsmosisQuery::PoolState { id: 2 }.into();
925        let state: PoolStateResponse = app.wrap().query(&query).unwrap();
926        let expected_assets = vec![coin(2_000_000 + 1993, "atom"), coin(1_000_000 - 993, "btc")];
927        assert_eq!(state.assets, expected_assets);
928    }
929
930    // TODO: make the following test work
931    #[test]
932    #[ignore]
933    fn estimate_swap_regression() {
934        let pool = Pool::new(coin(2_000_000, "atom"), coin(1_000_000, "btc"));
935
936        // set up with one pool
937        let mut app = OsmosisApp::new();
938        app.init_modules(|router, _, storage| {
939            router.custom.set_pool(storage, 1, &pool).unwrap();
940        });
941
942        // estimate the price (501505 * 0.997 = 500_000) after fees gone
943        let query = OsmosisQuery::estimate_swap(
944            MOCK_CONTRACT_ADDR,
945            1,
946            "atom",
947            "btc",
948            SwapAmount::In(Uint128::new(2007)),
949        );
950        let SwapResponse { amount } = app.wrap().query(&query.into()).unwrap();
951        // 6M * 1.5M = 2M * 4.5M -> output = 1.5M
952        let expected = SwapAmount::Out(Uint128::new(1000));
953        assert_eq!(amount, expected);
954
955        // now try the reverse query. we know what we need to pay to get 1.5M out
956        let query = OsmosisQuery::estimate_swap(
957            MOCK_CONTRACT_ADDR,
958            1,
959            "atom",
960            "btc",
961            SwapAmount::Out(Uint128::new(1000)),
962        );
963        let SwapResponse { amount } = app.wrap().query(&query.into()).unwrap();
964        let expected = SwapAmount::In(Uint128::new(2007));
965        assert_eq!(amount, expected);
966    }
967}