Merge pull request #111 from igneous-labs/feat/87-flat-fee-test

Address issues and test scaffolding for flat fee pricing program
igneous-labs · Dec 8, 2023 · 236474f · 236474f
2 parents 25f74f4 + f44fe1e
commit 236474f
Show file tree

Hide file tree

Showing 39 changed files with 309 additions and 154 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/generated/pricing-programs/flat_fee_interface/src/errors.rs b/generated/pricing-programs/flat_fee_interface/src/errors.rs
@@ -8,10 +8,14 @@ use thiserror::Error;
 pub enum FlatFeeError {
     #[error("Invalid program state data")]
     InvalidProgramStateData = 0,
+    #[error("Incorrect program state account")]
+    IncorrectProgramState = 1,
     #[error("FeeAccount is not initialized for the given LST mint")]
-    UnsupportedLstMint = 1,
+    UnsupportedLstMint = 2,
+    #[error("Given fee value is out of bound")]
+    SignedFeeOutOfBound = 3,
     #[error("Math error")]
-    MathError = 2,
+    MathError = 4,
 }
 impl From<FlatFeeError> for ProgramError {
     fn from(e: FlatFeeError) -> Self {

diff --git a/idl/pricing-programs/flat_fee.json b/idl/pricing-programs/flat_fee.json
@@ -410,11 +410,21 @@
     },
     {
       "code": 1,
+      "name": "IncorrectProgramState",
+      "msg": "Incorrect program state account"
+    },
+    {
+      "code": 2,
       "name": "UnsupportedLstMint",
       "msg": "FeeAccount is not initialized for the given LST mint"
     },
     {
-      "code": 2,
+      "code": 3,
+      "name": "SignedFeeOutOfBound",
+      "msg": "Given fee value is out of bound"
+    },
+    {
+      "code": 4,
       "name": "MathError",
       "msg": "Math error"
     }

diff --git a/libs/pricing-programs/flat-fee-lib/src/account_resolvers/add_lst.rs b/libs/pricing-programs/flat-fee-lib/src/account_resolvers/add_lst.rs
@@ -2,31 +2,49 @@ use flat_fee_interface::{AddLstKeys, FlatFeeError, ProgramState};
 use solana_program::{pubkey::Pubkey, system_program};
 use solana_readonly_account::{KeyedAccount, ReadonlyAccountData};
 
-use crate::{pda::FeeAccountFindPdaArgs, program, utils::try_program_state};
+use crate::{
+    pda::{FeeAccountCreatePdaArgs, FeeAccountFindPdaArgs},
+    program::STATE_ID,
+    utils::try_program_state,
+};
 
 pub struct AddLstFreeArgs<S: KeyedAccount + ReadonlyAccountData> {
     pub payer: Pubkey,
-    pub state: S,
+    pub state_acc: S,
     pub lst_mint: Pubkey,
 }
 
 impl<S: KeyedAccount + ReadonlyAccountData> AddLstFreeArgs<S> {
-    pub fn resolve(&self) -> Result<AddLstKeys, FlatFeeError> {
-        let bytes = &self.state.data();
+    pub fn resolve(self) -> Result<(AddLstKeys, FeeAccountCreatePdaArgs), FlatFeeError> {
+        let AddLstFreeArgs {
+            payer,
+            state_acc,
+            lst_mint,
+        } = self;
+
+        if *state_acc.key() != STATE_ID {
+            return Err(FlatFeeError::IncorrectProgramState);
+        }
+
+        let bytes = &state_acc.data();
         let state: &ProgramState = try_program_state(bytes)?;
 
-        let find_pda_args = FeeAccountFindPdaArgs {
-            lst_mint: self.lst_mint,
-        };
-        let (fee_acc, _bump) = find_pda_args.get_fee_account_address_and_bump_seed();
+        let find_pda_args = FeeAccountFindPdaArgs { lst_mint };
+        let (fee_acc, bump) = find_pda_args.get_fee_account_address_and_bump_seed();
 
-        Ok(AddLstKeys {
-            manager: state.manager,
-            payer: self.payer,
-            fee_acc,
-            lst_mint: self.lst_mint,
-            state: program::STATE_ID,
-            system_program: system_program::ID,
-        })
+        Ok((
+            AddLstKeys {
+                manager: state.manager,
+                payer,
+                fee_acc,
+                lst_mint,
+                state: STATE_ID,
+                system_program: system_program::ID,
+            },
+            FeeAccountCreatePdaArgs {
+                find_pda_args,
+                bump: [bump],
+            },
+        ))
     }
 }
diff --git a/libs/pricing-programs/flat-fee-lib/src/account_resolvers/remove_lst.rs b/libs/pricing-programs/flat-fee-lib/src/account_resolvers/remove_lst.rs
@@ -4,31 +4,41 @@ use solana_readonly_account::{KeyedAccount, ReadonlyAccountData};
 
 use crate::{
     pda::{FeeAccountCreatePdaArgs, FeeAccountFindPdaArgs},
-    program,
+    program::STATE_ID,
     utils::try_program_state,
 };
 
-pub struct RemoveLstWithMintFreeArgs<S: KeyedAccount + ReadonlyAccountData> {
+pub struct RemoveLstByMintFreeArgs<S: KeyedAccount + ReadonlyAccountData> {
     pub refund_rent_to: Pubkey,
     pub lst_mint: Pubkey,
-    pub state: S,
+    pub state_acc: S,
 }
 
-impl<S: KeyedAccount + ReadonlyAccountData> RemoveLstWithMintFreeArgs<S> {
-    fn resolve_with_fee_acc(&self, fee_acc: Pubkey) -> Result<RemoveLstKeys, FlatFeeError> {
-        let bytes = &self.state.data();
+impl<S: KeyedAccount + ReadonlyAccountData> RemoveLstByMintFreeArgs<S> {
+    fn resolve_with_fee_acc(self, fee_acc: Pubkey) -> Result<RemoveLstKeys, FlatFeeError> {
+        let RemoveLstByMintFreeArgs {
+            refund_rent_to,
+            lst_mint: _,
+            state_acc,
+        } = self;
+
+        if *state_acc.key() != STATE_ID {
+            return Err(FlatFeeError::IncorrectProgramState);
+        }
+
+        let bytes = &state_acc.data();
         let state: &ProgramState = try_program_state(bytes)?;
 
         Ok(RemoveLstKeys {
             manager: state.manager,
-            refund_rent_to: self.refund_rent_to,
+            refund_rent_to,
             fee_acc,
-            state: program::STATE_ID,
+            state: STATE_ID,
             system_program: system_program::ID,
         })
     }
 
-    pub fn resolve(&self) -> Result<RemoveLstKeys, FlatFeeError> {
+    pub fn resolve(self) -> Result<RemoveLstKeys, FlatFeeError> {
         let find_pda_args = FeeAccountFindPdaArgs {
             lst_mint: self.lst_mint,
         };
@@ -37,7 +47,7 @@ impl<S: KeyedAccount + ReadonlyAccountData> RemoveLstWithMintFreeArgs<S> {
         self.resolve_with_fee_acc(fee_acc)
     }
 
-    pub fn resolve_with_fee_acc_bump(&self, bump: u8) -> Result<RemoveLstKeys, FlatFeeError> {
+    pub fn resolve_with_fee_acc_bump(self, bump: u8) -> Result<RemoveLstKeys, FlatFeeError> {
         let create_pda_args = FeeAccountCreatePdaArgs {
             find_pda_args: FeeAccountFindPdaArgs {
                 lst_mint: self.lst_mint,
@@ -55,19 +65,28 @@ impl<S: KeyedAccount + ReadonlyAccountData> RemoveLstWithMintFreeArgs<S> {
 pub struct RemoveLstFreeArgs<S: KeyedAccount + ReadonlyAccountData> {
     pub refund_rent_to: Pubkey,
     pub fee_acc: Pubkey,
-    pub state: S,
+    pub state_acc: S,
 }
 
 impl<S: KeyedAccount + ReadonlyAccountData> RemoveLstFreeArgs<S> {
-    pub fn resolve(&self) -> Result<RemoveLstKeys, FlatFeeError> {
-        let bytes = &self.state.data();
+    pub fn resolve(self) -> Result<RemoveLstKeys, FlatFeeError> {
+        let RemoveLstFreeArgs {
+            refund_rent_to,
+            fee_acc,
+            state_acc,
+        } = self;
+        if *state_acc.key() != STATE_ID {
+            return Err(FlatFeeError::IncorrectProgramState);
+        }
+
+        let bytes = &state_acc.data();
         let state: &ProgramState = try_program_state(bytes)?;
 
         Ok(RemoveLstKeys {
             manager: state.manager,
-            refund_rent_to: self.refund_rent_to,
-            fee_acc: self.fee_acc,
-            state: program::STATE_ID,
+            refund_rent_to,
+            fee_acc,
+            state: STATE_ID,
             system_program: system_program::ID,
         })
     }

diff --git a/libs/pricing-programs/flat-fee-lib/src/account_resolvers/set_lp_withdrawal_fee.rs b/libs/pricing-programs/flat-fee-lib/src/account_resolvers/set_lp_withdrawal_fee.rs
@@ -1,7 +1,7 @@
 use flat_fee_interface::SetLpWithdrawalFeeKeys;
 use solana_program::pubkey::Pubkey;
 
-use crate::program;
+use crate::program::STATE_ID;
 
 pub struct SetLpWithdrawalFeeFreeArgs {
     pub manager: Pubkey,
@@ -11,7 +11,7 @@ impl SetLpWithdrawalFeeFreeArgs {
     pub fn resolve(self) -> SetLpWithdrawalFeeKeys {
         SetLpWithdrawalFeeKeys {
             manager: self.manager,
-            state: program::STATE_ID,
+            state: STATE_ID,
         }
     }
 }
diff --git a/libs/pricing-programs/flat-fee-lib/src/account_resolvers/set_lst_fee.rs b/libs/pricing-programs/flat-fee-lib/src/account_resolvers/set_lst_fee.rs
@@ -4,28 +4,37 @@ use solana_readonly_account::{KeyedAccount, ReadonlyAccountData};
 
 use crate::{
     pda::{FeeAccountCreatePdaArgs, FeeAccountFindPdaArgs},
-    program,
+    program::STATE_ID,
     utils::try_program_state,
 };
 
-pub struct SetLstFeeWithMintFreeArgs<S: KeyedAccount + ReadonlyAccountData> {
+pub struct SetLstFeeByMintFreeArgs<S: KeyedAccount + ReadonlyAccountData> {
     pub lst_mint: Pubkey,
-    pub state: S,
+    pub state_acc: S,
 }
 
-impl<S: KeyedAccount + ReadonlyAccountData> SetLstFeeWithMintFreeArgs<S> {
-    fn resolve_with_fee_acc(&self, fee_acc: Pubkey) -> Result<SetLstFeeKeys, FlatFeeError> {
-        let bytes = &self.state.data();
+impl<S: KeyedAccount + ReadonlyAccountData> SetLstFeeByMintFreeArgs<S> {
+    fn resolve_with_fee_acc(self, fee_acc: Pubkey) -> Result<SetLstFeeKeys, FlatFeeError> {
+        let SetLstFeeByMintFreeArgs {
+            lst_mint: _,
+            state_acc,
+        } = self;
+
+        if *state_acc.key() != STATE_ID {
+            return Err(FlatFeeError::IncorrectProgramState);
+        }
+
+        let bytes = &state_acc.data();
         let state: &ProgramState = try_program_state(bytes)?;
 
         Ok(SetLstFeeKeys {
             manager: state.manager,
             fee_acc,
-            state: program::STATE_ID,
+            state: STATE_ID,
         })
     }
 
-    pub fn resolve(&self) -> Result<SetLstFeeKeys, FlatFeeError> {
+    pub fn resolve(self) -> Result<SetLstFeeKeys, FlatFeeError> {
         let find_pda_args = FeeAccountFindPdaArgs {
             lst_mint: self.lst_mint,
         };
@@ -34,7 +43,7 @@ impl<S: KeyedAccount + ReadonlyAccountData> SetLstFeeWithMintFreeArgs<S> {
         self.resolve_with_fee_acc(fee_acc)
     }
 
-    pub fn resolve_with_fee_acc_bump(&self, bump: u8) -> Result<SetLstFeeKeys, FlatFeeError> {
+    pub fn resolve_with_fee_acc_bump(self, bump: u8) -> Result<SetLstFeeKeys, FlatFeeError> {
         let create_pda_args = FeeAccountCreatePdaArgs {
             find_pda_args: FeeAccountFindPdaArgs {
                 lst_mint: self.lst_mint,
@@ -51,18 +60,27 @@ impl<S: KeyedAccount + ReadonlyAccountData> SetLstFeeWithMintFreeArgs<S> {
 
 pub struct SetLstFeeFreeArgs<S: KeyedAccount + ReadonlyAccountData> {
     pub fee_acc: Pubkey,
-    pub state: S,
+    pub state_acc: S,
 }
 
 impl<S: KeyedAccount + ReadonlyAccountData> SetLstFeeFreeArgs<S> {
-    pub fn resolve(&self) -> Result<SetLstFeeKeys, FlatFeeError> {
-        let bytes = &self.state.data();
+    pub fn resolve(self) -> Result<SetLstFeeKeys, FlatFeeError> {
+        let SetLstFeeFreeArgs {
+            fee_acc: _,
+            state_acc,
+        } = self;
+
+        if *state_acc.key() != STATE_ID {
+            return Err(FlatFeeError::IncorrectProgramState);
+        }
+
+        let bytes = &state_acc.data();
         let state: &ProgramState = try_program_state(bytes)?;
 
         Ok(SetLstFeeKeys {
             manager: state.manager,
             fee_acc: self.fee_acc,
-            state: program::STATE_ID,
+            state: STATE_ID,
         })
     }
 }
diff --git a/libs/pricing-programs/flat-fee-lib/src/account_resolvers/set_manager.rs b/libs/pricing-programs/flat-fee-lib/src/account_resolvers/set_manager.rs
@@ -2,22 +2,31 @@ use flat_fee_interface::{FlatFeeError, ProgramState, SetManagerKeys};
 use solana_program::pubkey::Pubkey;
 use solana_readonly_account::{KeyedAccount, ReadonlyAccountData};
 
-use crate::{program, utils::try_program_state};
+use crate::{program::STATE_ID, utils::try_program_state};
 
 pub struct SetManagerFreeArgs<S: KeyedAccount + ReadonlyAccountData> {
     pub new_manager: Pubkey,
-    pub state: S,
+    pub state_acc: S,
 }
 
 impl<S: KeyedAccount + ReadonlyAccountData> SetManagerFreeArgs<S> {
     pub fn resolve(self) -> Result<SetManagerKeys, FlatFeeError> {
-        let bytes = &self.state.data();
+        let SetManagerFreeArgs {
+            new_manager,
+            state_acc,
+        } = self;
+
+        if *state_acc.key() != STATE_ID {
+            return Err(FlatFeeError::IncorrectProgramState);
+        }
+
+        let bytes = &state_acc.data();
         let state: &ProgramState = try_program_state(bytes)?;
 
         Ok(SetManagerKeys {
             current_manager: state.manager,
-            new_manager: self.new_manager,
-            state: program::STATE_ID,
+            new_manager,
+            state: STATE_ID,
         })
     }
 }
diff --git a/libs/pricing-programs/flat-fee-lib/src/calc/price_exact_in.rs b/libs/pricing-programs/flat-fee-lib/src/calc/price_exact_in.rs
@@ -3,10 +3,19 @@ use sanctum_token_ratio::{U64RatioFloor, BPS_DENOMINATOR};
 
 use super::BPS_DENOMINATOR_I16;
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct CalculatePriceExactInArgs {
+    pub input_fee_bps: i16,
+    pub output_fee_bps: i16,
+    pub sol_value: u64,
+}
+
 pub fn calculate_price_exact_in(
-    input_fee_bps: i16,
-    output_fee_bps: i16,
-    sol_value: u64,
+    CalculatePriceExactInArgs {
+        input_fee_bps,
+        output_fee_bps,
+        sol_value,
+    }: CalculatePriceExactInArgs,
 ) -> Result<u64, FlatFeeError> {
     let fee_bps = input_fee_bps
         .checked_add(output_fee_bps)

diff --git a/libs/pricing-programs/flat-fee-lib/src/calc/price_exact_out.rs b/libs/pricing-programs/flat-fee-lib/src/calc/price_exact_out.rs
@@ -3,10 +3,18 @@ use sanctum_token_ratio::{U64RatioFloor, BPS_DENOMINATOR};
 
 use super::BPS_DENOMINATOR_I16;
 
+pub struct CalculatePriceExactOut {
+    pub input_fee_bps: i16,
+    pub output_fee_bps: i16,
+    pub sol_value: u64,
+}
+
 pub fn calculate_price_exact_out(
-    input_fee_bps: i16,
-    output_fee_bps: i16,
-    sol_value: u64,
+    CalculatePriceExactOut {
+        input_fee_bps,
+        output_fee_bps,
+        sol_value,
+    }: CalculatePriceExactOut,
 ) -> Result<u64, FlatFeeError> {
     let fee_bps = input_fee_bps
         .checked_add(output_fee_bps)