[sui 5/x] - Sui clock, error codes, fee computation (#740)

* state getters and setters, change Move.toml dependency to sui/integration_v2

* finish state.move

* add new line to pyth

* use deployer cap pattern for state module

* sui pyth

* update price feeds, dynamic object fields, Sui object PriceInfoObject

* register price info object with pyth state after creation

* sui governance

* some newlines

* error codes

* update and comment

target_chains/sui/contracts/sources/batch_price_attestation.move

@@ -1,28 +1,24 @@
 module pyth::batch_price_attestation {
-
-    use sui::tx_context::{Self, TxContext};
+    use std::vector::{Self};
+    use sui::clock::{Self, Clock};
 
     use pyth::price_feed::{Self};
     use pyth::price_info::{Self, PriceInfo};
     use pyth::price_identifier::{Self};
     use pyth::price_status;
     use pyth::deserialize::{Self};
-    // TODO - Import Sui clock and use it for timekeeping instead of tx_context::epoch.
-    //        Replace epoch in deserialize_price_info with sui clock timestamp, and usage
-    //        of epoch in test_deserialize_batch_price_attestation.
-    // TODO - Use specific error messages in this module, specifically
-    //        for invalid_attestation_magic_value and invalid_batch_attestation_header_size.
+
     use wormhole::cursor::{Self, Cursor};
     use wormhole::bytes::{Self};
 
-    use std::vector::{Self};
-
     #[test_only]
     use pyth::price;
     #[test_only]
     use pyth::i64;
 
     const MAGIC: u64 = 0x50325748; // "P2WH" (Pyth2Wormhole) raw ASCII bytes
+    const E_INVALID_ATTESTATION_MAGIC_VALUE: u64 = 0;
+    const E_INVALID_BATCH_ATTESTATION_HEADER_SIZE: u64 = 1;
 
     struct BatchPriceAttestation {
         header: Header,
@@ -41,13 +37,13 @@ module pyth::batch_price_attestation {
 
     fun deserialize_header(cur: &mut Cursor<u8>): Header {
         let magic = (deserialize::deserialize_u32(cur) as u64);
-        assert!(magic == MAGIC, 0); // TODO - add specific error value - error::invalid_attestation_magic_value()
+        assert!(magic == MAGIC, E_INVALID_ATTESTATION_MAGIC_VALUE);
         let version_major = deserialize::deserialize_u16(cur);
         let version_minor = deserialize::deserialize_u16(cur);
         let header_size = deserialize::deserialize_u16(cur);
         let payload_id = deserialize::deserialize_u8(cur);
 
-        assert!(header_size >= 1, 0); // TODO - add specific error value - error::invalid_batch_attestation_header_size()
+        assert!(header_size >= 1, E_INVALID_BATCH_ATTESTATION_HEADER_SIZE);
         let unknown_header_bytes = header_size - 1;
         let _unknown = bytes::take_bytes(cur, (unknown_header_bytes as u64));
 
@@ -84,7 +80,7 @@ module pyth::batch_price_attestation {
         vector::borrow(&batch.price_infos, index)
     }
 
-    public fun deserialize(bytes: vector<u8>, ctx: &mut TxContext): BatchPriceAttestation {
+    public fun deserialize(bytes: vector<u8>, clock: &Clock): BatchPriceAttestation {
         let cur = cursor::new(bytes);
         let header = deserialize_header(&mut cur);
 
@@ -94,7 +90,7 @@ module pyth::batch_price_attestation {
 
         let i = 0;
         while (i < attestation_count) {
-            let price_info = deserialize_price_info(&mut cur, ctx);
+            let price_info = deserialize_price_info(&mut cur, clock);
             vector::push_back(&mut price_infos, price_info);
 
             // Consume any excess bytes
@@ -113,7 +109,7 @@ module pyth::batch_price_attestation {
         }
     }
 
-    fun deserialize_price_info(cur: &mut Cursor<u8>, ctx: &mut TxContext): PriceInfo {
+    fun deserialize_price_info(cur: &mut Cursor<u8>, clock: &Clock): PriceInfo {
 
         // Skip obselete field
         let _product_identifier = deserialize::deserialize_vector(cur, 32);
@@ -155,7 +151,7 @@ module pyth::batch_price_attestation {
 
         price_info::new_price_info(
             attestation_time,
-            tx_context::epoch(ctx), //TODO - use Sui Clock to get timestamp in seconds
+            clock::timestamp_ms(clock) / 1000, // Divide by 1000 to get timestamp in seconds
             price_feed::new(
                 price_identifier,
                 current_price,
@@ -167,21 +163,30 @@ module pyth::batch_price_attestation {
     #[test]
     #[expected_failure]
     fun test_deserialize_batch_price_attestation_invalid_magic() {
-        use sui::test_scenario::{Self, ctx};
+        use sui::test_scenario::{Self, take_shared, return_shared, ctx};
         let test = test_scenario::begin(@0x1234);
+        clock::create_for_testing(ctx(&mut test));
+        test_scenario::next_tx(&mut test, @0x1234);
+        let test_clock = take_shared<Clock>(&test);
 
         // A batch price attestation with a magic number of 0x50325749
         let bytes = x"5032574900030000000102000400951436e0be37536be96f0896366089506a59763d036728332d3e3038047851aea7c6c75c89f14810ec1c54c03ab8f1864a4c4032791f05747f560faec380a695d1000000000000049a0000000000000008fffffffb00000000000005dc0000000000000003000000000100000001000000006329c0eb000000006329c0e9000000006329c0e400000000000006150000000000000007215258d81468614f6b7e194c5d145609394f67b041e93e6695dcc616faadd0603b9551a68d01d954d6387aff4df1529027ffb2fee413082e509feb29cc4904fe000000000000041a0000000000000003fffffffb00000000000005cb0000000000000003010000000100000001000000006329c0eb000000006329c0e9000000006329c0e4000000000000048600000000000000078ac9cf3ab299af710d735163726fdae0db8465280502eb9f801f74b3c1bd190333832fad6e36eb05a8972fe5f219b27b5b2bb2230a79ce79beb4c5c5e7ecc76d00000000000003f20000000000000002fffffffb00000000000005e70000000000000003010000000100000001000000006329c0eb000000006329c0e9000000006329c0e40000000000000685000000000000000861db714e9ff987b6fedf00d01f9fea6db7c30632d6fc83b7bc9459d7192bc44a21a28b4c6619968bd8c20e95b0aaed7df2187fd310275347e0376a2cd7427db800000000000006cb0000000000000001fffffffb00000000000005e40000000000000003010000000100000001000000006329c0eb000000006329c0e9000000006329c0e400000000000007970000000000000001";
-        let _ = destroy(deserialize(bytes, ctx(&mut test)));
+        let _ = destroy(deserialize(bytes, &test_clock));
+        return_shared(test_clock);
         test_scenario::end(test);
     }
 
     #[test]
     fun test_deserialize_batch_price_attestation() {
-        use sui::test_scenario::{Self, ctx};
+        use sui::test_scenario::{Self, take_shared, return_shared, ctx};
         // Set the arrival time
         let test = test_scenario::begin(@0x1234);
-        let arrival_time = tx_context::epoch(ctx(&mut test));
+        clock::create_for_testing(ctx(&mut test));
+        test_scenario::next_tx(&mut test, @0x1234);
+        let test_clock = take_shared<Clock>(&test);
+        let arrival_time_in_seconds = clock::timestamp_ms(&test_clock) / 1000;
+
+        // let arrival_time = tx_context::epoch(ctx(&mut test));
 
         // A raw batch price attestation
         // The first attestation has a status of UNKNOWN
@@ -200,7 +205,7 @@ module pyth::batch_price_attestation {
             price_infos: vector<PriceInfo>[
                 price_info::new_price_info(
                     1663680747,
-                    arrival_time,
+                    arrival_time_in_seconds,
                     price_feed::new(
                         price_identifier::from_byte_vec(x"c6c75c89f14810ec1c54c03ab8f1864a4c4032791f05747f560faec380a695d1"),
                         price::new(i64::new(1557, false), 7, i64::new(5, true), 1663680740),
@@ -208,7 +213,7 @@ module pyth::batch_price_attestation {
                     )                ),
                 price_info::new_price_info(
                     1663680747,
-                    arrival_time,
+                    arrival_time_in_seconds,
                     price_feed::new(
                         price_identifier::from_byte_vec(x"3b9551a68d01d954d6387aff4df1529027ffb2fee413082e509feb29cc4904fe"),
                         price::new(i64::new(1050, false), 3, i64::new(5, true), 1663680745),
@@ -216,7 +221,7 @@ module pyth::batch_price_attestation {
                     )                ),
                 price_info::new_price_info(
                     1663680747,
-                    arrival_time,
+                    arrival_time_in_seconds,
                     price_feed::new(
                         price_identifier::from_byte_vec(x"33832fad6e36eb05a8972fe5f219b27b5b2bb2230a79ce79beb4c5c5e7ecc76d"),
                         price::new(i64::new(1010, false), 2, i64::new(5, true), 1663680745),
@@ -224,21 +229,22 @@ module pyth::batch_price_attestation {
                     )                ),
                 price_info::new_price_info(
                     1663680747,
-                    arrival_time,
+                    arrival_time_in_seconds,
                     price_feed::new(
                         price_identifier::from_byte_vec(x"21a28b4c6619968bd8c20e95b0aaed7df2187fd310275347e0376a2cd7427db8"),
                         price::new(i64::new(1739, false), 1, i64::new(5, true), 1663680745),
                         price::new(i64::new(1508, false), 3, i64::new(5, true), 1663680745),
-                    )                ),
+                    )
+                ),
             ],
         };
 
-        let deserialized = deserialize(bytes, ctx(&mut test));
+        let deserialized = deserialize(bytes, &test_clock);
 
         assert!(&expected == &deserialized, 1);
         destroy(expected);
         destroy(deserialized);
-
+        return_shared(test_clock);
         test_scenario::end(test);
     }
 }

target_chains/sui/contracts/sources/data_source.move

@@ -8,6 +8,8 @@ module pyth::data_source {
     use wormhole::external_address::ExternalAddress;
 
     const KEY: vector<u8> = b"data_sources";
+    const E_DATA_SOURCE_REGISTRY_ALREADY_EXISTS: u64 = 0;
+    const E_DATA_SOURCE_ALREADY_REGISTERED: u64 = 1;
 
     struct DataSource has copy, drop, store {
         emitter_chain: u64,
@@ -17,7 +19,7 @@ module pyth::data_source {
     public fun new_data_source_registry(parent_id: &mut UID, ctx: &mut TxContext) {
         assert!(
             !dynamic_field::exists_(parent_id, KEY),
-            0 // TODO - add custom error type
+            E_DATA_SOURCE_REGISTRY_ALREADY_EXISTS // TODO - add custom error type
         );
         dynamic_field::add(
             parent_id,
@@ -29,7 +31,7 @@ module pyth::data_source {
     public fun add(parent_id: &mut UID, data_source: DataSource) {
         assert!(
             !contains(parent_id, data_source),
-            0 // TODO - add custom error message
+            E_DATA_SOURCE_ALREADY_REGISTERED
         );
         set::add(
             dynamic_field::borrow_mut(parent_id, KEY),

target_chains/sui/contracts/sources/event.move

@@ -15,7 +15,7 @@ module pyth::event {
         timestamp: u64,
     }
 
-    public(friend) fun emit_price_feed_update(price_feed: PriceFeed, timestamp: u64) {
+    public(friend) fun emit_price_feed_update(price_feed: PriceFeed, timestamp: u64 /* in seconds */) {
         event::emit(
             PriceFeedUpdateEvent {
                 price_feed,

target_chains/sui/contracts/sources/governance/governance.move

@@ -15,6 +15,11 @@ module pyth::governance {
     use wormhole::vaa::{Self, VAA};
     use wormhole::state::{State as WormState};
 
+    const E_GOVERNANCE_CONTRACT_UPGRADE_CHAIN_ID_ZERO: u64 = 0;
+    const E_INVALID_GOVERNANCE_ACTION: u64 = 1;
+    const E_INVALID_GOVERNANCE_DATA_SOURCE: u64 = 2;
+    const E_INVALID_GOVERNANCE_SEQUENCE_NUMBER: u64 = 3;
+    
     public entry fun execute_governance_instruction(
         pyth_state : &mut State,
         worm_state: &WormState,
@@ -28,7 +33,7 @@ module pyth::governance {
         let action = governance_instruction::get_action(&instruction);
         if (action == governance_action::new_contract_upgrade()) {
             assert!(governance_instruction::get_target_chain_id(&instruction) != 0,
-                0); // TODO - error::governance_contract_upgrade_chain_id_zero()
+                E_GOVERNANCE_CONTRACT_UPGRADE_CHAIN_ID_ZERO);
             contract_upgrade::execute(worm_state, pyth_state, governance_instruction::destroy(instruction));
         } else if (action == governance_action::new_set_governance_data_source()) {
             set_governance_data_source::execute(pyth_state, governance_instruction::destroy(instruction));
@@ -40,7 +45,7 @@ module pyth::governance {
             set_stale_price_threshold::execute(pyth_state, governance_instruction::destroy(instruction));
         } else {
             governance_instruction::destroy(instruction);
-            assert!(false, 0); // TODO - error::invalid_governance_action()
+            assert!(false, E_INVALID_GOVERNANCE_ACTION);
         }
     }
 
@@ -59,11 +64,11 @@ module pyth::governance {
                 data_source::new(
                     (vaa::emitter_chain(&parsed_vaa) as u64),
                     vaa::emitter_address(&parsed_vaa))),
-            0); // TODO - error::invalid_governance_data_source()
+            E_INVALID_GOVERNANCE_DATA_SOURCE);
 
         // Check that the sequence number is greater than the last executed governance VAA
         let sequence = vaa::sequence(&parsed_vaa);
-        assert!(sequence > state::get_last_executed_governance_sequence(pyth_state), 0); // TODO - error::invalid_governance_sequence_number()
+        assert!(sequence > state::get_last_executed_governance_sequence(pyth_state), E_INVALID_GOVERNANCE_SEQUENCE_NUMBER);
         state::set_last_executed_governance_sequence(pyth_state, sequence);
 
         parsed_vaa

target_chains/sui/contracts/sources/governance/governance_action.move

@@ -6,13 +6,14 @@ module pyth::governance_action {
     const SET_DATA_SOURCES: u8 = 2;
     const SET_UPDATE_FEE: u8 = 3;
     const SET_STALE_PRICE_THRESHOLD: u8 = 4;
+    const E_INVALID_GOVERNANCE_ACTION: u64 = 5;
 
     struct GovernanceAction has copy, drop {
         value: u8,
     }
 
     public fun from_u8(value: u8): GovernanceAction {
-        assert!(CONTRACT_UPGRADE <= value && value <= SET_STALE_PRICE_THRESHOLD, 0); //TODO - add specific error: error::invalid_governance_action()
+        assert!(CONTRACT_UPGRADE <= value && value <= SET_STALE_PRICE_THRESHOLD, E_INVALID_GOVERNANCE_ACTION);
         GovernanceAction { value }
     }
 

target_chains/sui/contracts/sources/governance/governance_instruction.move

@@ -6,6 +6,10 @@ module pyth::governance_instruction {
     const MAGIC: vector<u8> = x"5054474d"; // "PTGM": Pyth Governance Message
     const MODULE: u8 = 1;
 
+    const E_INVALID_GOVERNANCE_MODULE: u64 = 0;
+    const E_INVALID_GOVERNANCE_MAGIC_VALUE: u64 = 1;
+    const E_TARGET_CHAIN_MISMATCH: u64 = 2;
+    
     struct GovernanceInstruction {
         module_: u8,
         action: GovernanceAction,
@@ -14,15 +18,16 @@ module pyth::governance_instruction {
     }
 
     fun validate(instruction: &GovernanceInstruction) {
-        assert!(instruction.module_ == MODULE, 0); // TODO - add custom error::invalid_governance_module()
+        assert!(instruction.module_ == MODULE, E_INVALID_GOVERNANCE_MODULE);
         let target_chain_id = instruction.target_chain_id;
-        assert!(target_chain_id == (wormhole::state::chain_id() as u64) || target_chain_id == 0, 0); // TODO - custom error: error::invalid_governance_target_chain_id()
+        assert!(target_chain_id == (wormhole::state::chain_id() as u64) || target_chain_id == 0, E_TARGET_CHAIN_MISMATCH);
     }
 
     public fun from_byte_vec(bytes: vector<u8>): GovernanceInstruction {
         let cursor = cursor::new(bytes);
         let magic = deserialize::deserialize_vector(&mut cursor, 4);
-        assert!(magic == MAGIC, 0); // TODO error::invalid_governance_magic_value()
+        assert!(magic == MAGIC, E_INVALID_GOVERNANCE_MAGIC_VALUE);
+        // "module" is a reserved keyword, so we use "module_" instead.
         let module_ = deserialize::deserialize_u8(&mut cursor);
         let action = governance_action::from_u8(deserialize::deserialize_u8(&mut cursor));
         let target_chain_id = deserialize::deserialize_u16(&mut cursor);

target_chains/sui/contracts/sources/governance/set_update_fee.move

@@ -10,6 +10,7 @@ module pyth::set_update_fee {
     friend pyth::governance;
 
     const MAX_U64: u128 = (1 << 64) - 1;
+    const E_EXPONENT_DOES_NOT_FIT_IN_U8: u64 = 0;
 
     struct SetUpdateFee {
         mantissa: u64,
@@ -18,7 +19,7 @@ module pyth::set_update_fee {
 
     public(friend) fun execute(pyth_state: &mut State, payload: vector<u8>) {
         let SetUpdateFee { mantissa, exponent } = from_byte_vec(payload);
-        assert!(exponent <= 255, 0); // TODO - throw error that exponent does not fit in a u8
+        assert!(exponent <= 255, E_EXPONENT_DOES_NOT_FIT_IN_U8);
         let fee = apply_exponent(mantissa, (exponent as u8));
         state::set_base_update_fee(pyth_state, fee);
     }

target_chains/sui/contracts/sources/price_info.move

@@ -8,6 +8,8 @@ module pyth::price_info {
     use pyth::price_identifier::{PriceIdentifier};
 
     const KEY: vector<u8> = b"price_info";
+    const E_PRICE_INFO_REGISTRY_ALREADY_EXISTS: u64 = 0;
+    const E_PRICE_IDENTIFIER_ALREADY_REGISTERED: u64 = 1;
 
     friend pyth::pyth;
 
@@ -30,7 +32,7 @@ module pyth::price_info {
     public fun new_price_info_registry(parent_id: &mut UID, ctx: &mut TxContext) {
         assert!(
             !dynamic_object_field::exists_(parent_id, KEY),
-            0 // TODO - add custom error message
+            E_PRICE_INFO_REGISTRY_ALREADY_EXISTS
         );
         dynamic_object_field::add(
             parent_id,
@@ -42,7 +44,7 @@ module pyth::price_info {
     public fun add(parent_id: &mut UID, price_identifier: PriceIdentifier, id: ID) {
         assert!(
             !contains(parent_id, price_identifier),
-            0 // TODO - add custom error message
+            E_PRICE_IDENTIFIER_ALREADY_REGISTERED
         );
         table::add(
             dynamic_object_field::borrow_mut(parent_id, KEY),

target_chains/sui/contracts/sources/pyth.move

@@ -1,10 +1,10 @@
 module pyth::pyth {
     use std::vector;
     use sui::tx_context::{TxContext};
-    use sui::coin::{Coin};
+    use sui::coin::{Self, Coin};
     use sui::sui::{SUI};
     use sui::transfer::{Self};
-    use sui::tx_context::{Self};
+    use sui::clock::{Self, Clock};
 
     use pyth::event::{Self as pyth_event};
     use pyth::data_source::{Self, DataSource};
@@ -18,6 +18,9 @@ module pyth::pyth {
     use wormhole::vaa::{Self};
     use wormhole::state::{State as WormState};
 
+    const E_DATA_SOURCE_EMITTER_ADDRESS_AND_CHAIN_IDS_DIFFERENT_LENGTHS: u64 = 0;
+    const E_INVALID_DATA_SOURCE: u64 = 1;
+    const E_INSUFFICIENT_FEE: u64 = 2;
 
     /// Call init_and_share_state with deployer cap to initialize
     /// state and emit event corresponding to Pyth initialization.
@@ -54,8 +57,8 @@ module pyth::pyth {
         emitter_addresses: vector<vector<u8>>
     ): vector<DataSource> {
 
-        // TODO - add custom error type error::data_source_emitter_address_and_chain_ids_different_lengths()
-        assert!(vector::length(&emitter_chain_ids) == vector::length(&emitter_addresses), 0);
+        assert!(vector::length(&emitter_chain_ids) == vector::length(&emitter_addresses),
+            E_DATA_SOURCE_EMITTER_ADDRESS_AND_CHAIN_IDS_DIFFERENT_LENGTHS);
 
         let sources = vector::empty();
         let i = 0;
@@ -75,6 +78,7 @@ module pyth::pyth {
         worm_state: &WormState,
         pyth_state: &mut PythState,
         vaas: vector<vector<u8>>,
+        clock: &Clock,
         ctx: &mut TxContext
     ){
         while (!vector::is_empty(&vaas)) {
@@ -91,10 +95,10 @@ module pyth::pyth {
                         (vaa::emitter_chain(&vaa) as u64),
                         vaa::emitter_address(&vaa))
                     ),
-            0); // TODO - use custom error message - error::invalid_data_source()
+            E_INVALID_DATA_SOURCE);
 
             // Deserialize the batch price attestation
-            let price_infos = batch_price_attestation::destroy(batch_price_attestation::deserialize(vaa::take_payload(vaa), ctx));
+            let price_infos = batch_price_attestation::destroy(batch_price_attestation::deserialize(vaa::take_payload(vaa), clock));
             while (!vector::is_empty(&price_infos)){
                 let cur_price_info = vector::pop_back(&mut price_infos);
 
@@ -140,11 +144,13 @@ module pyth::pyth {
         vaas: vector<vector<u8>>,
         price_info_objects: &mut vector<PriceInfoObject>,
         fee: Coin<SUI>,
+        clock: &Clock,
         ctx: &mut TxContext
     ) {
         // Charge the message update fee
-        // TODO - error::insufficient_fee()
-        //assert!(get_update_fee(&vaas) <= coin::value(&fee), 0);
+        assert!(get_total_update_fee(pyth_state, &vaas) <= coin::value(&fee), E_INSUFFICIENT_FEE);
+
+        // TODO: use Wormhole fee collector instead of transferring funds to deployer address.
         transfer::public_transfer(fee, @pyth);
 
         // Update the price feed from each VAA
@@ -154,6 +160,7 @@ module pyth::pyth {
                 pyth_state,
                 vector::pop_back(&mut vaas),
                 price_info_objects,
+                clock,
                 ctx
             );
         };
@@ -167,6 +174,7 @@ module pyth::pyth {
         pyth_state: &PythState,
         worm_vaa: vector<u8>,
         price_info_objects: &mut vector<PriceInfoObject>,
+        clock: &Clock,
         ctx: &mut TxContext
     ) {
         // Deserialize the VAA
@@ -180,20 +188,20 @@ module pyth::pyth {
                     (vaa::emitter_chain(&vaa) as u64),
                     vaa::emitter_address(&vaa))
                 ),
-        0); // TODO - use custom error message - error::invalid_data_source()
+        E_INVALID_DATA_SOURCE);
 
         // Deserialize the batch price attestation
-        let price_infos = batch_price_attestation::destroy(batch_price_attestation::deserialize(vaa::take_payload(vaa), ctx));
+        let price_infos = batch_price_attestation::destroy(batch_price_attestation::deserialize(vaa::take_payload(vaa), clock));
 
         // Update price info objects.
-        update_cache(price_infos, price_info_objects, ctx);
+        update_cache(price_infos, price_info_objects, clock);
     }
 
     /// Update PriceInfoObjects using up-to-date PriceInfos.
     fun update_cache(
         updates: vector<PriceInfo>,
         price_info_objects: &mut vector<PriceInfoObject>,
-        ctx: &mut TxContext
+        clock: &Clock,
     ){
         while (!vector::is_empty(&updates)) {
             let update = vector::pop_back(&mut updates);
@@ -209,8 +217,7 @@ module pyth::pyth {
                 if (price_info::get_price_identifier(&price_info) ==
                     price_info::get_price_identifier(&update)){
                     found = true;
-                    // TODO: use clock timestamp instead of epoch in the future
-                    pyth_event::emit_price_feed_update(price_feed::from(price_info::get_price_feed(&update)), tx_context::epoch(ctx));
+                    pyth_event::emit_price_feed_update(price_feed::from(price_info::get_price_feed(&update)), clock::timestamp_ms(clock)/1000);
 
                     // Update the price info object with the new updated price info.
                     if (is_fresh_update(&update, vector::borrow(price_info_objects, i))){
@@ -243,4 +250,14 @@ module pyth::pyth {
 
         update_timestamp > cached_timestamp
     }
+
+    /// Get the number of AptosCoin's required to perform the given price updates.
+    ///
+    /// Please read more information about the update fee here: https://docs.pyth.network/consume-data/on-demand#fees
+    public fun get_total_update_fee(pyth_state: &PythState, update_data: &vector<vector<u8>>): u64 {
+        state::get_base_update_fee(pyth_state) * vector::length(update_data)
+    }
 }
+
+// TODO - pyth tests
+// https://github.com/pyth-network/pyth-crosschain/blob/main/target_chains/aptos/contracts/sources/pyth.move#L384