From 176cc5a3a415e8cb61df7ccd31df2c92dded8b7c Mon Sep 17 00:00:00 2001
From: 0xripleys <0xripleys@solend.fi>
Date: Thu, 31 Oct 2024 20:45:45 -0400
Subject: [PATCH] add passthrough functions

---
 contracts/sources/fees.move         |  1 -
 contracts/sources/hooks/weight.move | 26 ++++++++++++++++++++++++++
 contracts/tests/weight_tests.move   | 16 ++++++++++++++++
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/contracts/sources/fees.move b/contracts/sources/fees.move
index 24de118..8999772 100644
--- a/contracts/sources/fees.move
+++ b/contracts/sources/fees.move
@@ -1,6 +1,5 @@
 module liquid_staking::fees {
     use sui::bag::{Self, Bag};
-    use std::u64::max;
 
     // Errors
     const EInvalidFeeConfig: u64 = 0;
diff --git a/contracts/sources/hooks/weight.move b/contracts/sources/hooks/weight.move
index 231c5ab..6303f3f 100644
--- a/contracts/sources/hooks/weight.move
+++ b/contracts/sources/hooks/weight.move
@@ -3,10 +3,13 @@
 module liquid_staking::weight {
     use sui_system::sui_system::{SuiSystemState};
     use liquid_staking::liquid_staking::{LiquidStakingInfo, AdminCap};
+    use liquid_staking::fees::{FeeConfig};
     use sui::vec_map::{Self, VecMap};
     use sui::bag::{Self, Bag};
     use liquid_staking::version::{Self, Version};
     use sui::package;
+    use sui::coin::Coin;
+    use sui::sui::SUI;
 
     /* Constants */
     const CURRENT_VERSION: u16 = 1;
@@ -64,6 +67,29 @@ module liquid_staking::weight {
         self.total_weight = total_weight;
     }
 
+    public fun update_fees<P>(
+        self: &mut WeightHook<P>,
+        _: &WeightHookAdminCap<P>,
+        liquid_staking_info: &mut LiquidStakingInfo<P>,
+        fee_config: FeeConfig,
+    ) {
+        self.version.assert_version_and_upgrade(CURRENT_VERSION);
+
+        liquid_staking_info.update_fees(&self.admin_cap, fee_config);
+    }
+
+    public fun collect_fees<P>(
+        self: &mut WeightHook<P>,
+        _: &WeightHookAdminCap<P>,
+        liquid_staking_info: &mut LiquidStakingInfo<P>,
+        system_state: &mut SuiSystemState,
+        ctx: &mut TxContext,
+    ): Coin<SUI> {
+        self.version.assert_version_and_upgrade(CURRENT_VERSION);
+
+        liquid_staking_info.collect_fees(system_state, &self.admin_cap, ctx)
+    }
+
     public fun rebalance<P>(
         self: &mut WeightHook<P>,
         system_state: &mut SuiSystemState,
diff --git a/contracts/tests/weight_tests.move b/contracts/tests/weight_tests.move
index 2c70667..3b2d6df 100644
--- a/contracts/tests/weight_tests.move
+++ b/contracts/tests/weight_tests.move
@@ -103,6 +103,20 @@ module liquid_staking::weight_tests {
         assert!(lst_info.storage().validators().borrow(1).total_sui_amount() == 0, 0);
         assert!(lst_info.storage().validators().borrow(2).total_sui_amount() == 100 * MIST_PER_SUI, 0);
 
+        // test update fees
+        let new_fees = fees::new_builder(scenario.ctx()).set_sui_mint_fee_bps(100).to_fee_config();
+        weight_hook.update_fees(&weight_hook_admin_cap, &mut lst_info, new_fees);
+
+        assert!(lst_info.fee_config().sui_mint_fee_bps() == 100, 0);
+
+        // mint some lst
+        let sui = coin::mint_for_testing(100 * MIST_PER_SUI, scenario.ctx());
+        let lst2 = lst_info.mint(&mut system_state, sui, scenario.ctx());
+
+        // test collect fees
+        let collected_fees = weight_hook.collect_fees(&weight_hook_admin_cap, &mut lst_info, &mut system_state, scenario.ctx());
+        assert!(collected_fees.value() == MIST_PER_SUI, 0);
+
         // sharing to make sure shared object deletion actually works lol
         transfer::public_share_object(weight_hook);
         scenario.next_tx(@0x0);
@@ -115,6 +129,8 @@ module liquid_staking::weight_tests {
         sui::test_utils::destroy(admin_cap);
         sui::test_utils::destroy(lst_info);
         sui::test_utils::destroy(lst);
+        sui::test_utils::destroy(lst2);
+        sui::test_utils::destroy(collected_fees);
 
         scenario.end(); 
      }