From 9d29e8d0136e9f1042452f40b5f71ad1b4afe81d Mon Sep 17 00:00:00 2001
From: lemunozm <lemunozm@gmail.com>
Date: Thu, 21 Sep 2023 15:52:36 +0200
Subject: [PATCH] remove portfolio entry when removing a loan. add portfolio
 tests

---
 pallets/loans/src/lib.rs             |  4 ++
 pallets/loans/src/types/portfolio.rs | 59 ++++++++++++++++++++++++++--
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/pallets/loans/src/lib.rs b/pallets/loans/src/lib.rs
index b76bed756d..4ff5063295 100644
--- a/pallets/loans/src/lib.rs
+++ b/pallets/loans/src/lib.rs
@@ -947,6 +947,10 @@ pub mod pallet {
 					.position(|(id, _)| *id == loan_id)
 					.ok_or(Error::<T>::LoanNotActiveOrNotFound)?;
 
+				PortfolioValuation::<T>::try_mutate(pool_id, |portfolio| {
+					portfolio.remove_elem(loan_id)
+				})?;
+
 				Ok((
 					active_loans.swap_remove(index).1,
 					active_loans.len().ensure_into()?,
diff --git a/pallets/loans/src/types/portfolio.rs b/pallets/loans/src/types/portfolio.rs
index 1dfebe8055..65fe54102d 100644
--- a/pallets/loans/src/types/portfolio.rs
+++ b/pallets/loans/src/types/portfolio.rs
@@ -71,11 +71,11 @@ where
 		self.last_updated
 	}
 
-	pub fn value_of(&self, id: ElemId) -> Option<&Balance> {
+	pub fn value_of(&self, id: ElemId) -> Option<Balance> {
 		self.values
 			.iter()
 			.find(|(elem_id, _)| *elem_id == id)
-			.map(|(_, balance)| balance)
+			.map(|(_, balance)| *balance)
 	}
 
 	pub fn insert_elem(&mut self, id: ElemId, pv: Balance) -> DispatchResult {
@@ -83,7 +83,8 @@ where
 			.try_push((id, pv))
 			.map_err(|_| DispatchError::Other("Max portfolio size reached"))?;
 
-		Ok(self.value.ensure_add_assign(pv)?)
+		self.value.ensure_add_assign(pv)?;
+		Ok(())
 	}
 
 	pub fn update_elem(&mut self, id: ElemId, new_pv: Balance) -> DispatchResult {
@@ -110,6 +111,18 @@ where
 
 		Ok(())
 	}
+
+	pub fn remove_elem(&mut self, elem_id: ElemId) -> DispatchResult {
+		let index = self
+			.values
+			.iter()
+			.position(|(id, _)| *id == elem_id)
+			.ok_or(DispatchError::CannotLookup)?;
+
+		let (_, pv) = self.values.swap_remove(index);
+		self.value.ensure_sub_assign(pv)?;
+		Ok(())
+	}
 }
 
 /// Type that builds a PortfolioValuation with the current instant.
@@ -136,3 +149,43 @@ pub enum PortfolioValuationUpdateType {
 	/// Portfolio Valuation was updated inexactly based on loan status changes
 	Inexact,
 }
+
+#[cfg(test)]
+mod tests {
+	use frame_support::assert_ok;
+	use sp_core::ConstU32;
+
+	use super::*;
+
+	#[test]
+	fn general_usage() {
+		let mut portfolio = PortfolioValuation::<u128, u64, ConstU32<3>>::new(10);
+
+		assert_ok!(portfolio.insert_elem(1, 100));
+		assert_ok!(portfolio.insert_elem(2, 200));
+		assert_ok!(portfolio.insert_elem(3, 300));
+
+		assert_eq!(portfolio.value(), 600);
+
+		// Increase
+		assert_ok!(portfolio.update_elem(1, 300));
+		assert_eq!(portfolio.value(), 800);
+
+		// Do not change
+		assert_ok!(portfolio.update_elem(2, 200));
+		assert_eq!(portfolio.value(), 800);
+
+		// Decrease
+		assert_ok!(portfolio.update_elem(3, 100));
+		assert_eq!(portfolio.value(), 600);
+
+		assert_eq!(portfolio.value_of(1), Some(300));
+		assert_eq!(portfolio.value_of(2), Some(200));
+		assert_eq!(portfolio.value_of(3), Some(100));
+
+		assert_ok!(portfolio.remove_elem(1));
+		assert_ok!(portfolio.remove_elem(2));
+		assert_ok!(portfolio.remove_elem(3));
+		assert_eq!(portfolio.value(), 0);
+	}
+}