Merge pull request #635 from galacticcouncil/fix/stablewap-tradable-s…

…tate fix: check if pool and asset exists when setting tradable state [B3]
galacticcouncil · Jul 7, 2023 · fa71c82 · fa71c82
2 parents be86148 + 4cb5f08
commit fa71c82
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 3 deletions.
diff --git a/pallets/stableswap/Cargo.toml b/pallets/stableswap/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = 'pallet-stableswap'
-version = '2.0.1'
+version = '2.0.2'
 description = 'AMM for correlated assets'
 authors = ['GalacticCouncil']
 edition = '2021'

diff --git a/pallets/stableswap/src/lib.rs b/pallets/stableswap/src/lib.rs
@@ -717,6 +717,9 @@ pub mod pallet {
 		) -> DispatchResult {
 			T::AuthorityOrigin::ensure_origin(origin)?;
 
+			let pool = Pools::<T>::get(pool_id).ok_or(Error::<T>::PoolNotFound)?;
+			let _ = pool.find_asset(asset_id).ok_or(Error::<T>::AssetNotInPool)?;
+
 			AssetTradability::<T>::mutate(pool_id, asset_id, |current_state| {
 				*current_state = state;
 			});

diff --git a/pallets/stableswap/src/tests/update_pool.rs b/pallets/stableswap/src/tests/update_pool.rs
@@ -1,6 +1,6 @@
 use crate::tests::mock::*;
-use crate::types::PoolInfo;
-use crate::{Error, Pools};
+use crate::types::{PoolInfo, Tradability};
+use crate::{AssetTradability, Error, Pools};
 use frame_support::{assert_noop, assert_ok};
 use sp_runtime::Permill;
 use std::num::NonZeroU16;
@@ -242,3 +242,78 @@ fn update_pool_should_fail_when_pool_does_not_exists() {
 			);
 		});
 }
+
+#[test]
+fn set_tradable_state_should_work_when_asset_in_pool() {
+	let asset_a: AssetId = 1;
+	let asset_b: AssetId = 2;
+	let pool_id: AssetId = 100;
+
+	ExtBuilder::default()
+		.with_endowed_accounts(vec![(ALICE, asset_a, 200 * ONE), (ALICE, asset_b, 200 * ONE)])
+		.with_registered_asset("pool".as_bytes().to_vec(), pool_id)
+		.with_registered_asset("one".as_bytes().to_vec(), asset_a)
+		.with_registered_asset("two".as_bytes().to_vec(), asset_b)
+		.build()
+		.execute_with(|| {
+			assert_ok!(Stableswap::create_pool(
+				RuntimeOrigin::root(),
+				pool_id,
+				vec![asset_a, asset_b],
+				100,
+				Permill::from_percent(0),
+				Permill::from_percent(0),
+			));
+
+			assert_ok!(Stableswap::set_asset_tradable_state(
+				RuntimeOrigin::root(),
+				pool_id,
+				asset_a,
+				Tradability::FROZEN,
+			));
+
+			assert_eq!(<AssetTradability<Test>>::get(pool_id, asset_a), Tradability::FROZEN,);
+		});
+}
+#[test]
+fn set_tradable_state_should_fail_when_asset_not_in_pool() {
+	let asset_a: AssetId = 1;
+	let asset_b: AssetId = 2;
+	let pool_id: AssetId = 100;
+
+	ExtBuilder::default()
+		.with_endowed_accounts(vec![(ALICE, asset_a, 200 * ONE), (ALICE, asset_b, 200 * ONE)])
+		.with_registered_asset("pool".as_bytes().to_vec(), pool_id)
+		.with_registered_asset("one".as_bytes().to_vec(), asset_a)
+		.with_registered_asset("two".as_bytes().to_vec(), asset_b)
+		.build()
+		.execute_with(|| {
+			assert_ok!(Stableswap::create_pool(
+				RuntimeOrigin::root(),
+				pool_id,
+				vec![asset_a, asset_b],
+				100,
+				Permill::from_percent(0),
+				Permill::from_percent(0),
+			));
+
+			assert_noop!(
+				Stableswap::set_asset_tradable_state(RuntimeOrigin::root(), pool_id, 3, Tradability::FROZEN,),
+				Error::<Test>::AssetNotInPool
+			);
+		});
+}
+#[test]
+fn set_tradable_state_should_fail_when_pool_does_not_exist() {
+	let pool_id: AssetId = 100;
+
+	ExtBuilder::default()
+		.with_registered_asset("pool".as_bytes().to_vec(), pool_id)
+		.build()
+		.execute_with(|| {
+			assert_noop!(
+				Stableswap::set_asset_tradable_state(RuntimeOrigin::root(), pool_id, 1, Tradability::FROZEN,),
+				Error::<Test>::PoolNotFound
+			);
+		});
+}