Fix tests

anza-xyz · Oct 25, 2024 · 29e3e59 · 29e3e59
1 parent 220cc98
commit 29e3e59
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 15 deletions.
diff --git a/compute-budget/src/compute_budget_limits.rs b/compute-budget/src/compute_budget_limits.rs
@@ -39,15 +39,20 @@ impl Default for ComputeBudgetLimits {
         }
     }
 }
+
+//#[cfg(test)]
+pub fn get_prioritization_fee_for_test(compute_unit_price: u64, compute_unit_limit: u64) -> u64 {
+    get_prioritization_fee(compute_unit_price, compute_unit_limit)
+}
+
 fn get_prioritization_fee(compute_unit_price: u64, compute_unit_limit: u64) -> u64 {
     let micro_lamport_fee: MicroLamports =
         (compute_unit_price as u128).saturating_mul(compute_unit_limit as u128);
-    let fee = micro_lamport_fee
+    micro_lamport_fee
         .saturating_add(MICRO_LAMPORTS_PER_LAMPORT.saturating_sub(1) as u128)
         .checked_div(MICRO_LAMPORTS_PER_LAMPORT as u128)
         .and_then(|fee| u64::try_from(fee).ok())
-        .unwrap_or(u64::MAX);
-    fee
+        .unwrap_or(u64::MAX)
 }
 
 impl From<ComputeBudgetLimits> for FeeBudgetLimits {

diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs
@@ -35,8 +35,7 @@ use {
     },
     solana_compute_budget::{
         compute_budget::ComputeBudget,
-        compute_budget_limits::{self, MAX_COMPUTE_UNIT_LIMIT},
-        prioritization_fee::{PrioritizationFeeDetails, PrioritizationFeeType},
+        compute_budget_limits::{self, MAX_COMPUTE_UNIT_LIMIT, get_prioritization_fee_for_test},
     },
     solana_feature_set::{self as feature_set, FeatureSet},
     solana_inline_spl::token,
@@ -10251,10 +10250,6 @@ fn test_calculate_fee_compute_units() {
         MAX_COMPUTE_UNIT_LIMIT,
     ] {
         const PRIORITIZATION_FEE_RATE: u64 = 42;
-        let prioritization_fee_details = PrioritizationFeeDetails::new(
-            PrioritizationFeeType::ComputeUnitPrice(PRIORITIZATION_FEE_RATE),
-            requested_compute_units as u64,
-        );
         let message = new_sanitized_message(Message::new(
             &[
                 ComputeBudgetInstruction::set_compute_unit_limit(requested_compute_units),
@@ -10266,7 +10261,7 @@ fn test_calculate_fee_compute_units() {
         let fee = calculate_test_fee(&message, 1, &fee_structure);
         assert_eq!(
             fee,
-            lamports_per_signature + prioritization_fee_details.get_fee()
+            lamports_per_signature + get_prioritization_fee_for_test(PRIORITIZATION_FEE_RATE, requested_compute_units as u64)
         );
     }
 }
@@ -10280,11 +10275,7 @@ fn test_calculate_prioritization_fee() {
 
     let request_units = 1_000_000_u32;
     let request_unit_price = 2_000_000_000_u64;
-    let prioritization_fee_details = PrioritizationFeeDetails::new(
-        PrioritizationFeeType::ComputeUnitPrice(request_unit_price),
-        request_units as u64,
-    );
-    let prioritization_fee = prioritization_fee_details.get_fee();
+    let prioritization_fee = get_prioritization_fee_for_test(request_unit_price, request_units as u64);
 
     let message = new_sanitized_message(Message::new(
         &[