From fcf0acbaa93b631a3618af07e986c5c9c1c4ed38 Mon Sep 17 00:00:00 2001
From: yooml <ymlll0508@gmail.com>
Date: Thu, 28 Dec 2023 22:07:39 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=F0=9F=90=9B=20reduce=5Famount=20(#1118)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pallets/leverage-staking/src/lib.rs   | 12 ++++++++----
 pallets/leverage-staking/src/tests.rs | 17 +++++++++++++++++
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/pallets/leverage-staking/src/lib.rs b/pallets/leverage-staking/src/lib.rs
index 8a2422131..cc85be842 100644
--- a/pallets/leverage-staking/src/lib.rs
+++ b/pallets/leverage-staking/src/lib.rs
@@ -146,10 +146,14 @@ impl<T: Config> Pallet<T> {
 
 		match rate.cmp(&current_rate) {
 			Ordering::Less => {
-				let reduce_amount = current_rate
-					.checked_sub(&rate)
-					.and_then(|r| r.checked_mul_int(base_token_value))
-					.ok_or(ArithmeticError::Overflow)?;
+				let reduce_amount = if rate.is_zero() {
+					account_borrows
+				} else {
+					current_rate
+						.checked_sub(&rate)
+						.and_then(|r| r.checked_mul_int(base_token_value))
+						.ok_or(ArithmeticError::Overflow)?
+				};
 				Self::reduce_leverage(&who, asset_id, vtoken_id, reduce_amount)?;
 			},
 			Ordering::Equal => return Err(Error::<T>::ArgumentsError.into()),
diff --git a/pallets/leverage-staking/src/tests.rs b/pallets/leverage-staking/src/tests.rs
index ac4f4a237..b46230852 100644
--- a/pallets/leverage-staking/src/tests.rs
+++ b/pallets/leverage-staking/src/tests.rs
@@ -141,12 +141,14 @@ fn reduce_leverage_should_work() {
 			Deposits { voucher_balance: 9500000, is_collateral: true },
 		);
 		assert_eq!(Tokens::balance(VDOT, &1), 9999999900000);
+		assert_eq!(Tokens::balance(DOT, &1), 990000000000000);
 		assert_ok!(LeverageStaking::flash_loan_deposit(
 			RuntimeOrigin::signed(1),
 			DOT,
 			FixedU128::from_inner(unit(800_000)),
 		));
 		assert_eq!(Tokens::balance(VDOT, &1), 9999999900000);
+		assert_eq!(Tokens::balance(DOT, &1), 990000000000098);
 		assert_eq!(
 			AccountBorrows::<Test>::get(DOT, 1),
 			BorrowSnapshot { principal: 80_000, borrow_index: 1.into() },
@@ -155,5 +157,20 @@ fn reduce_leverage_should_work() {
 			AccountDeposits::<Test>::get(VDOT, 1),
 			Deposits { voucher_balance: 8994050, is_collateral: true },
 		);
+		assert_ok!(LeverageStaking::flash_loan_deposit(
+			RuntimeOrigin::signed(1),
+			DOT,
+			FixedU128::from_inner(0),
+		));
+		assert_eq!(Tokens::balance(VDOT, &1), 9999999900000);
+		assert_eq!(Tokens::balance(DOT, &1), 990000000000196);
+		assert_eq!(
+			AccountBorrows::<Test>::get(DOT, 1),
+			BorrowSnapshot { principal: 0, borrow_index: 1.into() },
+		);
+		assert_eq!(
+			AccountDeposits::<Test>::get(VDOT, 1),
+			Deposits { voucher_balance: 4981100, is_collateral: true },
+		);
 	});
 }