From 3d76021802601c44e7844da45d660d62a16dd23d Mon Sep 17 00:00:00 2001
From: Ayush Suresh <byteslayer31037@gmail.com>
Date: Wed, 6 Aug 2025 16:27:10 -0500
Subject: [PATCH] fixed order

---
 .../pyth-receiver/src/integration_tests.rs    | 29 +++++++++++++++++++
 .../stylus/contracts/pyth-receiver/src/lib.rs | 10 +++----
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/target_chains/stylus/contracts/pyth-receiver/src/integration_tests.rs b/target_chains/stylus/contracts/pyth-receiver/src/integration_tests.rs
index 99d8c9ca0d..6e63657f68 100644
--- a/target_chains/stylus/contracts/pyth-receiver/src/integration_tests.rs
+++ b/target_chains/stylus/contracts/pyth-receiver/src/integration_tests.rs
@@ -8,6 +8,7 @@ mod test {
     use motsu::prelude::*;
     use pythnet_sdk::wire::v1::{AccumulatorUpdateData, Proof};
     use std::time::Duration;
+    use stylus_sdk::types::AddressVM;
     use wormhole_contract::WormholeContract;
 
     const PYTHNET_CHAIN_ID: u16 = 26;
@@ -106,8 +107,12 @@ mod test {
         let result = pyth_contract
             .sender_and_value(alice, update_fee)
             .update_price_feeds(update_data);
+
         assert!(result.is_ok());
 
+        assert_eq!(alice.balance(), U256::ZERO);
+        assert_eq!(pyth_contract.balance(), update_fee);
+
         let price_result = pyth_contract
             .sender(alice)
             .get_price_unsafe(ban_usd_feed_id());
@@ -157,11 +162,17 @@ mod test {
             .update_price_feeds(update_data1);
         assert!(result1.is_ok());
 
+        assert_eq!(alice.balance(), update_fee2);
+        assert_eq!(pyth_contract.balance(), update_fee1);
+
         let result2 = pyth_contract
             .sender_and_value(alice, update_fee2)
             .update_price_feeds(update_data2);
         assert!(result2.is_ok());
 
+        assert_eq!(alice.balance(), U256::ZERO);
+        assert_eq!(pyth_contract.balance(), update_fee1 + update_fee2);
+
         let price_result = pyth_contract
             .sender(alice)
             .get_price_unsafe(ban_usd_feed_id());
@@ -231,6 +242,9 @@ mod test {
             .update_price_feeds(update_data);
         assert!(result.is_ok());
 
+        assert_eq!(alice.balance(), U256::ZERO);
+        assert_eq!(pyth_contract.balance(), update_fee);
+
         let price_result = pyth_contract
             .sender(alice)
             .get_price_no_older_than(btc_usd_feed_id(), u64::MAX);
@@ -257,6 +271,9 @@ mod test {
             .update_price_feeds(update_data);
         assert!(result.is_ok());
 
+        assert_eq!(alice.balance(), U256::ZERO);
+        assert_eq!(pyth_contract.balance(), update_fee);
+
         let price_result = pyth_contract
             .sender(alice)
             .get_price_no_older_than(btc_usd_feed_id(), 1);
@@ -286,6 +303,9 @@ mod test {
             .update_price_feeds(update_data);
         assert!(result.is_ok());
 
+        assert_eq!(alice.balance(), U256::ZERO);
+        assert_eq!(pyth_contract.balance(), update_fee);
+
         let first_price_result = pyth_contract
             .sender(alice)
             .get_price_unsafe(ban_usd_feed_id());
@@ -327,6 +347,9 @@ mod test {
             .update_price_feeds(update_data);
         assert!(result.is_ok());
 
+        assert_eq!(alice.balance(), U256::ZERO);
+        assert_eq!(pyth_contract.balance(), update_fee);
+
         assert!(pyth_contract
             .sender(alice)
             .price_feed_exists(ban_usd_feed_id()));
@@ -368,6 +391,9 @@ mod test {
             .sender_and_value(alice, update_fee)
             .update_price_feeds(update_data);
 
+        assert_eq!(alice.balance(), U256::ZERO);
+        assert_eq!(pyth_contract.balance(), update_fee);
+
         assert!(result.is_ok());
 
         let price_result = pyth_contract
@@ -395,6 +421,9 @@ mod test {
             .sender_and_value(alice, update_fee)
             .update_price_feeds(update_data);
 
+        assert_eq!(alice.balance(), U256::ZERO);
+        assert_eq!(pyth_contract.balance(), update_fee);
+
         assert!(result.is_ok());
 
         let price_result1 = pyth_contract
diff --git a/target_chains/stylus/contracts/pyth-receiver/src/lib.rs b/target_chains/stylus/contracts/pyth-receiver/src/lib.rs
index ab88aa1428..58e541e9c9 100644
--- a/target_chains/stylus/contracts/pyth-receiver/src/lib.rs
+++ b/target_chains/stylus/contracts/pyth-receiver/src/lib.rs
@@ -213,17 +213,17 @@ impl PythReceiver {
         &mut self,
         update_data: Vec<Vec<u8>>,
     ) -> Result<(), PythReceiverError> {
-        for data in &update_data {
-            self.update_price_feeds_internal(data.clone(), 0, 0, false)?;
-        }
-
-        let total_fee = self.get_update_fee(update_data)?;
+        let total_fee = self.get_update_fee(update_data.clone())?;
 
         let value = self.vm().msg_value();
 
         if value < total_fee {
             return Err(PythReceiverError::InsufficientFee);
         }
+
+        for data in &update_data {
+            self.update_price_feeds_internal(data.clone(), 0, 0, false)?;
+        }
         Ok(())
     }