[sui 6/x] - Pyth unit tests (#757)

* 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

* unit tests for pyth.move, add UpgradeCap to Pyth State (will be used for contract upgrades)

* updates

* test_get_update_fee test passes

* fix test_get_update_fee and test_update_price_feeds_corrupt_vaa

* test_update_price_feeds_invalid_data_source

* test_create_and_update_price_feeds

* test_create_and_update_price_feeds_success and test_create_and_update_price_feeds_price_info_object_not_found_failure

* test_update_cache

* update

* test_update_cache_old_update

* update_price_feeds_if_fresh

* comment

target_chains/sui/contracts/Move.lock

@@ -0,0 +1,29 @@
+# @generated by Move, please check-in and do not edit manually.
+
+[move]
+version = 0
+
+dependencies = [
+  { name = "Sui" },
+  { name = "Wormhole" },
+]
+
+[[move.package]]
+name = "MoveStdlib"
+source = { git = "https://github.com/MystenLabs/sui.git", rev = "81dbcf2b6cab07d623a1012bf31daf658963c765", subdir = "crates/sui-framework/packages/move-stdlib" }
+
+[[move.package]]
+name = "Sui"
+source = { git = "https://github.com/MystenLabs/sui.git", rev = "81dbcf2b6cab07d623a1012bf31daf658963c765", subdir = "crates/sui-framework/packages/sui-framework" }
+
+dependencies = [
+  { name = "MoveStdlib" },
+]
+
+[[move.package]]
+name = "Wormhole"
+source = { git = "https://github.com/wormhole-foundation/wormhole.git", rev = "sui/integration_v2", subdir = "sui/wormhole" }
+
+dependencies = [
+  { name = "Sui" },
+]

target_chains/sui/contracts/Move.toml

@@ -5,7 +5,7 @@ version = "0.0.1"
 [dependencies.Sui]
 git = "https://github.com/MystenLabs/sui.git"
 subdir = "crates/sui-framework/packages/sui-framework"
-rev = "82c9c80c11488858f1d3930f47ec9f335a566683"
+rev = "ddfc3fa0768a38286787319603a5458a9ff91cc1"
 
 [dependencies.Wormhole]
 git = "https://github.com/wormhole-foundation/wormhole.git"

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

@@ -1,5 +1,5 @@
 module pyth::governance {
-    use sui::tx_context::{TxContext};
+    use sui::clock::{Clock};
 
     use pyth::data_source::{Self};
     use pyth::governance_instruction;
@@ -19,14 +19,14 @@ module pyth::governance {
     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,
         vaa_bytes: vector<u8>,
-        ctx: &mut TxContext
+        clock: &Clock
     ) {
-        let parsed_vaa = parse_and_verify_governance_vaa(pyth_state, worm_state, vaa_bytes, ctx);
+        let parsed_vaa = parse_and_verify_governance_vaa(pyth_state, worm_state, vaa_bytes, clock);
         let instruction = governance_instruction::from_byte_vec(vaa::take_payload(parsed_vaa));
 
         // Dispatch the instruction to the appropiate handler
@@ -53,9 +53,9 @@ module pyth::governance {
         pyth_state: &mut State,
         worm_state: &WormState,
         bytes: vector<u8>,
-        ctx: &mut TxContext
+        clock: &Clock,
     ): VAA {
-        let parsed_vaa = vaa::parse_and_verify(worm_state, bytes, ctx);
+        let parsed_vaa = vaa::parse_and_verify(worm_state, bytes, clock);
 
         // Check that the governance data source is valid
         assert!(

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

@@ -6,6 +6,7 @@ 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 {

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

@@ -9,7 +9,7 @@ module pyth::governance_instruction {
     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,

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

@@ -3,6 +3,7 @@ module pyth::set_data_sources {
 
     use wormhole::cursor;
     use wormhole::external_address::{Self};
+    use wormhole::bytes32::{Self};
 
     use pyth::deserialize;
     use pyth::data_source::{Self, DataSource};
@@ -28,7 +29,7 @@ module pyth::set_data_sources {
         let i = 0;
         while (i < data_sources_count) {
             let emitter_chain_id = deserialize::deserialize_u16(&mut cursor);
-            let emitter_address = external_address::from_bytes(deserialize::deserialize_vector(&mut cursor, 32));
+            let emitter_address = external_address::new(bytes32::from_bytes(deserialize::deserialize_vector(&mut cursor, 32)));
             vector::push_back(&mut sources, data_source::new((emitter_chain_id as u64), emitter_address));
 
             i = i + 1;

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

@@ -5,6 +5,7 @@ module pyth::set_governance_data_source {
 
     use wormhole::cursor;
     use wormhole::external_address::{Self, ExternalAddress};
+    use wormhole::bytes32::{Self};
     //use wormhole::state::{Self}
 
     friend pyth::governance;
@@ -24,7 +25,7 @@ module pyth::set_governance_data_source {
     fun from_byte_vec(bytes: vector<u8>): SetGovernanceDataSource {
         let cursor = cursor::new(bytes);
         let emitter_chain_id = deserialize::deserialize_u16(&mut cursor);
-        let emitter_address = external_address::from_bytes(deserialize::deserialize_vector(&mut cursor, 32));
+        let emitter_address = external_address::new(bytes32::from_bytes(deserialize::deserialize_vector(&mut cursor, 32)));
         let initial_sequence = deserialize::deserialize_u64(&mut cursor);
         cursor::destroy_empty(cursor);
         SetGovernanceDataSource {

target_chains/sui/contracts/sources/price_info.move

@@ -13,8 +13,8 @@ module pyth::price_info {
 
     friend pyth::pyth;
 
-    /// Sui Object version of PriceInfo.
-    /// Has a key and lives in global store.
+    /// Sui object version of PriceInfo.
+    /// Has a key ability, is unique for each price identifier, and lives in global store.
     struct PriceInfoObject has key, store {
         id: UID,
         price_info: PriceInfo

target_chains/sui/contracts/sources/pyth.move

@@ -5,6 +5,7 @@ module pyth::pyth {
     use sui::sui::{SUI};
     use sui::transfer::{Self};
     use sui::clock::{Self, Clock};
+    use sui::package::{UpgradeCap};
 
     use pyth::event::{Self as pyth_event};
     use pyth::data_source::{Self, DataSource};
@@ -12,20 +13,29 @@ module pyth::pyth {
     use pyth::price_info::{Self, PriceInfo, PriceInfoObject};
     use pyth::batch_price_attestation::{Self};
     use pyth::price_feed::{Self};
-    use pyth::price::{Self};
+    use pyth::price::{Self, Price};
+    use pyth::price_identifier::{PriceIdentifier};
 
     use wormhole::external_address::{Self};
     use wormhole::vaa::{Self};
     use wormhole::state::{State as WormState};
+    use wormhole::bytes32::{Self};
 
     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;
+    const E_STALE_PRICE_UPDATE: u64 = 3;
+    const E_PRICE_INFO_OBJECT_NOT_FOUND: u64 = 4;
+    const E_INVALID_PUBLISH_TIMES_LENGTH: u64 = 5;
+    const E_NO_FRESH_DATA: u64 = 6;
+
+    friend pyth::pyth_tests;
 
     /// Call init_and_share_state with deployer cap to initialize
     /// state and emit event corresponding to Pyth initialization.
     public entry fun init_pyth(
         deployer: DeployerCap,
+        upgrade_cap: UpgradeCap,
         stale_price_threshold: u64,
         governance_emitter_chain_id: u64,
         governance_emitter_address: vector<u8>,
@@ -36,11 +46,13 @@ module pyth::pyth {
     ) {
         state::init_and_share_state(
             deployer,
+            upgrade_cap,
             stale_price_threshold,
             update_fee,
             data_source::new(
                 governance_emitter_chain_id,
-                external_address::from_bytes(governance_emitter_address)),
+                external_address::new((bytes32::from_bytes(governance_emitter_address)))
+            ),
             parse_data_sources(
                 data_sources_emitter_chain_ids,
                 data_sources_emitter_addresses,
@@ -65,7 +77,7 @@ module pyth::pyth {
         while (i < vector::length(&emitter_chain_ids)) {
             vector::push_back(&mut sources, data_source::new(
                 *vector::borrow(&emitter_chain_ids, i),
-                external_address::from_bytes(*vector::borrow(&emitter_addresses, i))
+                external_address::new(bytes32::from_bytes(*vector::borrow(&emitter_addresses, i)))
             ));
 
             i = i + 1;
@@ -85,7 +97,7 @@ module pyth::pyth {
             let vaa = vector::pop_back(&mut vaas);
 
             // Deserialize the VAA
-            let vaa = vaa::parse_and_verify(worm_state, vaa, ctx);
+            let vaa = vaa::parse_and_verify(worm_state, vaa, clock);
 
             // Check that the VAA is from a valid data source (emitter)
             assert!(
@@ -130,7 +142,7 @@ module pyth::pyth {
     ///
     /// The vaas argument is a vector of VAAs encoded as bytes.
     ///
-    /// The javascript https://github.com/pyth-network/pyth-js/tree/main/pyth-aptos-js package
+    /// The javascript https://github.com/pyth-network/pyth-js/tree/main/pyth-sui-js package
     /// should be used to fetch these VAAs from the Price Service. More information about this
     /// process can be found at https://docs.pyth.network/consume-data.
     ///
@@ -144,9 +156,8 @@ module pyth::pyth {
         vaas: vector<vector<u8>>,
         price_info_objects: &mut vector<PriceInfoObject>,
         fee: Coin<SUI>,
-        clock: &Clock,
-        ctx: &mut TxContext
-    ) {
+        clock: &Clock
+    ){
         // Charge the message update fee
         assert!(get_total_update_fee(pyth_state, &vaas) <= coin::value(&fee), E_INSUFFICIENT_FEE);
 
@@ -160,8 +171,7 @@ module pyth::pyth {
                 pyth_state,
                 vector::pop_back(&mut vaas),
                 price_info_objects,
-                clock,
-                ctx
+                clock
             );
         };
     }
@@ -174,11 +184,10 @@ module pyth::pyth {
         pyth_state: &PythState,
         worm_vaa: vector<u8>,
         price_info_objects: &mut vector<PriceInfoObject>,
-        clock: &Clock,
-        ctx: &mut TxContext
+        clock: &Clock
     ) {
         // Deserialize the VAA
-        let vaa = vaa::parse_and_verify(worm_state, worm_vaa, ctx);
+        let vaa = vaa::parse_and_verify(worm_state, worm_vaa, clock);
 
         // Check that the VAA is from a valid data source (emitter)
         assert!(
@@ -198,7 +207,7 @@ module pyth::pyth {
     }
 
     /// Update PriceInfoObjects using up-to-date PriceInfos.
-    fun update_cache(
+    public(friend) fun update_cache(
         updates: vector<PriceInfo>,
         price_info_objects: &mut vector<PriceInfoObject>,
         clock: &Clock,
@@ -208,9 +217,10 @@ module pyth::pyth {
             let i = 0;
             let found = false;
             // Find PriceInfoObjects corresponding to the current update (PriceInfo).
-            // TODO - This for loop might be expensive if there are a large
-            //        number of updates and/or price_info_objects we are updating.
-            while (i < vector::length<PriceInfoObject>(price_info_objects)){
+            // TODO - Construct an in-memory table to make look-ups faster?
+            //        This loop might be expensive if there are a large number
+            //        of updates and/or price_info_objects we are updating.
+            while (i < vector::length<PriceInfoObject>(price_info_objects) && found == false){
                 // Check if the current price info object corresponds to the price feed that
                 // the update is meant for.
                 let price_info = price_info::get_price_info_from_price_info_object(vector::borrow(price_info_objects, i));
@@ -226,16 +236,58 @@ module pyth::pyth {
                             update
                         );
                     }
-                }
+                };
+                i = i + 1;
             };
             if (!found){
-                // TODO - throw error, since the price_feeds in price_info_objects do
-                //        not constitute a superset of the price_feeds to be updated
+                abort(E_PRICE_INFO_OBJECT_NOT_FOUND)
             }
         };
         vector::destroy_empty(updates);
     }
 
+    /// Update the cached price feeds with the data in the given VAAs, using
+    /// update_price_feeds(). However, this function will only have an effect if any of the
+    /// prices in the update are fresh. The price_identifiers and publish_times parameters
+    /// are used to determine if the update is fresh without doing any serialisation or verification
+    /// of the VAAs, potentially saving time and gas. If the update contains no fresh data, this function
+    /// will revert with error::no_fresh_data().
+    ///
+    /// For a given price update i in the batch, that price is considered fresh if the current cached
+    /// price for price_identifiers[i] is older than publish_times[i].
+    public fun update_price_feeds_if_fresh(
+        vaas: vector<vector<u8>>,
+        worm_state: &WormState,
+        pyth_state: &PythState,
+        price_info_objects: &mut vector<PriceInfoObject>,
+        publish_times: vector<u64>,
+        fee: Coin<SUI>,
+        clock: &Clock
+    ) {
+        assert!(vector::length(price_info_objects) == vector::length(&publish_times),
+            E_INVALID_PUBLISH_TIMES_LENGTH
+        );
+
+        let fresh_data = false;
+        let i = 0;
+        while (i < vector::length(&publish_times)) {
+            let cur_price_info = price_info::get_price_info_from_price_info_object(vector::borrow(price_info_objects, i));
+            let cur_price_feed = price_info::get_price_feed(&cur_price_info);
+            let cur_price = price_feed::get_price(cur_price_feed);
+
+            let cached_timestamp = price::get_timestamp(&cur_price);
+            if (cached_timestamp < *vector::borrow(&publish_times, i)) {
+                fresh_data = true;
+                break
+            };
+
+            i = i + 1;
+        };
+
+        assert!(fresh_data, E_NO_FRESH_DATA);
+        update_price_feeds(worm_state, pyth_state, vaas, price_info_objects, fee, clock);
+    }
+
     /// Determine if the given price update is "fresh": we have nothing newer already cached for that
     /// price feed within a PriceInfoObject.
     fun is_fresh_update(update: &PriceInfo, price_info_object: &PriceInfoObject): bool {
@@ -251,13 +303,731 @@ module pyth::pyth {
         update_timestamp > cached_timestamp
     }
 
-    /// Get the number of AptosCoin's required to perform the given price updates.
+    // -----------------------------------------------------------------------------
+    // Query the cached prices
+    //
+    // It is strongly recommended to update the cached prices using the functions above,
+    // before using the functions below to query the cached data.
+
+    /// Determine if a price feed for the given price_identifier exists
+    public fun price_feed_exists(state: &PythState, price_identifier: PriceIdentifier): bool {
+        state::price_feed_object_exists(state, price_identifier)
+    }
+
+    /// Get the latest available price cached for the given price identifier, if that price is
+    /// no older than the stale price threshold.
+    ///
+    /// Please refer to the documentation at https://docs.pyth.network/consumers/best-practices for
+    /// how to how this price safely.
+    ///
+    /// Important: Pyth uses an on-demand update model, where consumers need to update the
+    /// cached prices before using them. Please read more about this at https://docs.pyth.network/consume-data/on-demand.
+    /// get_price() is likely to abort unless you call update_price_feeds() to update the cached price
+    /// beforehand, as the cached prices may be older than the stale price threshold.
+    ///
+    /// The price_info_object is a Sui object with the key ability that uniquely
+    /// contains a price feed for a given price_identifier.
+    ///
+    public fun get_price(state: &PythState, price_info_object: &PriceInfoObject, clock: &Clock): Price {
+        get_price_no_older_than(price_info_object, clock, state::get_stale_price_threshold_secs(state))
+    }
+
+    /// Get the latest available price cached for the given price identifier, if that price is
+    /// no older than the given age.
+    public fun get_price_no_older_than(price_info_object: &PriceInfoObject, clock: &Clock, max_age_secs: u64): Price {
+        let price = get_price_unsafe(price_info_object);
+        check_price_is_fresh(&price, clock, max_age_secs);
+        price
+    }
+
+    /// Get the latest available price cached for the given price identifier.
     ///
+    /// WARNING: the returned price can be from arbitrarily far in the past.
+    /// This function makes no guarantees that the returned price is recent or
+    /// useful for any particular application. Users of this function should check
+    /// the returned timestamp to ensure that the returned price is sufficiently
+    /// recent for their application. The checked get_price_no_older_than()
+    /// function should be used in preference to this.
+    public fun get_price_unsafe(price_info_object: &PriceInfoObject): Price {
+        // TODO: extract Price from this guy...
+        let price_info = price_info::get_price_info_from_price_info_object(price_info_object);
+        price_feed::get_price(
+            price_info::get_price_feed(&price_info)
+        )
+    }
+
+    fun abs_diff(x: u64, y: u64): u64 {
+        if (x > y) {
+            return x - y
+        } else {
+            return y - x
+        }
+    }
+
+    /// Get the stale price threshold: the amount of time after which a cached price
+    /// is considered stale and no longer returned by get_price()/get_ema_price().
+    public fun get_stale_price_threshold_secs(state: &PythState): u64 {
+        state::get_stale_price_threshold_secs(state)
+    }
+
+    fun check_price_is_fresh(price: &Price, clock: &Clock, max_age_secs: u64) {
+        let age = abs_diff(clock::timestamp_ms(clock)/1000, price::get_timestamp(price));
+        assert!(age < max_age_secs, E_STALE_PRICE_UPDATE);
+    }
+
     /// 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)
     }
 }
 
+module pyth::pyth_tests{
+    use std::vector::{Self};
+
+    use sui::sui::SUI;
+    use sui::coin::{Self, Coin};
+    use sui::clock::{Self, Clock};
+    use sui::test_scenario::{Self, Scenario, ctx, take_shared, return_shared};
+    use sui::package::Self;
+    use sui::object::{Self, ID};
+
+    use pyth::state::{Self, State as PythState};
+    use pyth::price_identifier::{Self};
+    use pyth::price_info::{Self, PriceInfo, PriceInfoObject};
+    use pyth::price_feed::{Self};
+    use pyth::data_source::{Self, DataSource};
+    use pyth::i64::{Self};
+    use pyth::price::{Self};
+    use pyth::pyth::{Self};
+
+    use wormhole::setup::{Self as wormhole_setup, DeployerCap};
+    use wormhole::external_address::{Self};
+    use wormhole::bytes32::{Self};
+    use wormhole::state::{State as WormState};
+
+    const DEPLOYER: address = @0x1234;
+
+    #[test_only]
+    /// A vector containing a single VAA with:
+    /// - emitter chain ID 17
+    /// - emitter address 0x71f8dcb863d176e2c420ad6610cf687359612b6fb392e0642b0ca6b1f186aa3b
+    /// - payload corresponding to the batch price attestation of the prices returned by get_mock_price_infos()
+    const TEST_VAAS: vector<vector<u8>> = vector[x"0100000000010036eb563b80a24f4253bee6150eb8924e4bdf6e4fa1dfc759a6664d2e865b4b134651a7b021b7f1ce3bd078070b688b6f2e37ce2de0d9b48e6a78684561e49d5201527e4f9b00000001001171f8dcb863d176e2c420ad6610cf687359612b6fb392e0642b0ca6b1f186aa3b0000000000000001005032574800030000000102000400951436e0be37536be96f0896366089506a59763d036728332d3e3038047851aea7c6c75c89f14810ec1c54c03ab8f1864a4c4032791f05747f560faec380a695d1000000000000049a0000000000000008fffffffb00000000000005dc0000000000000003000000000100000001000000006329c0eb000000006329c0e9000000006329c0e400000000000006150000000000000007215258d81468614f6b7e194c5d145609394f67b041e93e6695dcc616faadd0603b9551a68d01d954d6387aff4df1529027ffb2fee413082e509feb29cc4904fe000000000000041a0000000000000003fffffffb00000000000005cb0000000000000003010000000100000001000000006329c0eb000000006329c0e9000000006329c0e4000000000000048600000000000000078ac9cf3ab299af710d735163726fdae0db8465280502eb9f801f74b3c1bd190333832fad6e36eb05a8972fe5f219b27b5b2bb2230a79ce79beb4c5c5e7ecc76d00000000000003f20000000000000002fffffffb00000000000005e70000000000000003010000000100000001000000006329c0eb000000006329c0e9000000006329c0e40000000000000685000000000000000861db714e9ff987b6fedf00d01f9fea6db7c30632d6fc83b7bc9459d7192bc44a21a28b4c6619968bd8c20e95b0aaed7df2187fd310275347e0376a2cd7427db800000000000006cb0000000000000001fffffffb00000000000005e40000000000000003010000000100000001000000006329c0eb000000006329c0e9000000006329c0e400000000000007970000000000000001"];
+
+    #[test_only]
+    /// Init Wormhole core bridge state.
+    /// Init Pyth state.
+    /// Set initial Sui clock time.
+    /// Mint some SUI fee coins.
+    fun setup_test(
+        stale_price_threshold: u64,
+        governance_emitter_chain_id: u64,
+        governance_emitter_address: vector<u8>,
+        data_sources: vector<DataSource>,
+        base_update_fee: u64,
+        to_mint: u64
+    ): (Scenario, Coin<SUI>) {
+
+        let scenario = test_scenario::begin(DEPLOYER);
+
+        // Initialize Wormhole core bridge.
+        wormhole_setup::init_test_only(ctx(&mut scenario));
+        test_scenario::next_tx(&mut scenario, DEPLOYER);
+        // Take the `DeployerCap` from the sender of the transaction.
+        let deployer_cap =
+            test_scenario::take_from_address<DeployerCap>(
+                &scenario,
+                DEPLOYER
+            );
+
+        // This will be created and sent to the transaction sender automatically
+        // when the contract is published. This exists in place of grabbing
+        // it from the sender.
+        let upgrade_cap =
+            package::test_publish(
+                object::id_from_address(@wormhole),
+                test_scenario::ctx(&mut scenario)
+            );
+
+        let governance_chain = 1234;
+        let governance_contract =
+            x"deadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef";
+        let initial_guardians =
+            vector[
+                x"beFA429d57cD18b7F8A4d91A2da9AB4AF05d0FBe"
+            ];
+        let guardian_set_seconds_to_live = 5678;
+        let message_fee = 350;
+
+        wormhole_setup::complete(
+            deployer_cap,
+            upgrade_cap,
+            governance_chain,
+            governance_contract,
+            initial_guardians,
+            guardian_set_seconds_to_live,
+            message_fee,
+            test_scenario::ctx(&mut scenario)
+        );
+
+        // Create and share a global clock object for timekeeping.
+        clock::create_for_testing(ctx(&mut scenario));
+
+        // Initialize Pyth state.
+        let pyth_upgrade_cap=
+            package::test_publish(
+                object::id_from_address(@pyth),
+                test_scenario::ctx(&mut scenario)
+            );
+
+        state::init_test_only(ctx(&mut scenario));
+        test_scenario::next_tx(&mut scenario, DEPLOYER);
+        let pyth_deployer_cap = test_scenario::take_from_address<state::DeployerCap>(
+            &scenario,
+            DEPLOYER
+        );
+
+        state::init_and_share_state(
+            pyth_deployer_cap,
+            pyth_upgrade_cap,
+            stale_price_threshold,
+            base_update_fee,
+            data_source::new(governance_emitter_chain_id, external_address::new(bytes32::from_bytes(governance_emitter_address))),
+            data_sources,
+            ctx(&mut scenario)
+        );
+
+        let coins = coin::mint_for_testing<SUI>(to_mint, ctx(&mut scenario));
+        (scenario, coins)
+    }
+
+    #[test_only]
+    fun get_mock_price_infos(): vector<PriceInfo> {
+        vector<PriceInfo>[
+                price_info::new_price_info(
+                    1663680747,
+                    1663074349,
+                    price_feed::new(
+                        price_identifier::from_byte_vec(x"c6c75c89f14810ec1c54c03ab8f1864a4c4032791f05747f560faec380a695d1"),
+                        price::new(i64::new(1557, false), 7, i64::new(5, true), 1663680740),
+                        price::new(i64::new(1500, false), 3, i64::new(5, true), 1663680740),
+                    ),
+                ),
+                price_info::new_price_info(
+                    1663680747,
+                    1663074349,
+                    price_feed::new(
+                        price_identifier::from_byte_vec(x"3b9551a68d01d954d6387aff4df1529027ffb2fee413082e509feb29cc4904fe"),
+                        price::new(i64::new(1050, false), 3, i64::new(5, true), 1663680745),
+                        price::new(i64::new(1483, false), 3, i64::new(5, true), 1663680745),
+                    ),
+                ),
+                price_info::new_price_info(
+                    1663680747,
+                    1663074349,
+                    price_feed::new(
+                        price_identifier::from_byte_vec(x"33832fad6e36eb05a8972fe5f219b27b5b2bb2230a79ce79beb4c5c5e7ecc76d"),
+                        price::new(i64::new(1010, false), 2, i64::new(5, true), 1663680745),
+                        price::new(i64::new(1511, false), 3, i64::new(5, true), 1663680745),
+                    ),
+                ),
+                price_info::new_price_info(
+                    1663680747,
+                    1663074349,
+                    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),
+                    ),
+                ),
+            ]
+    }
+
+    #[test_only]
+    /// Compare the expected price feed with the actual Pyth price feeds.
+    fun check_price_feeds_cached(expected: &vector<PriceInfo>, actual: &vector<PriceInfoObject>) {
+        // Check that we can retrieve the correct current price and ema price for each price feed
+        let i = 0;
+        while (i < vector::length(expected)) {
+            let price_feed = price_info::get_price_feed(vector::borrow(expected, i));
+            let price = price_feed::get_price(price_feed);
+            let ema_price = price_feed::get_ema_price(price_feed);
+            let price_identifier = price_info::get_price_identifier(vector::borrow(expected, i));
+
+            let actual_price_info = price_info::get_price_info_from_price_info_object(vector::borrow(actual, i));
+            let actual_price_feed = price_info::get_price_feed(&actual_price_info);
+            let actual_price = price_feed::get_price(actual_price_feed);
+            let actual_ema_price = price_feed::get_ema_price(actual_price_feed);
+            let actual_price_identifier = price_info::get_price_identifier(&actual_price_info);
+
+            assert!(price == actual_price, 0);
+            assert!(ema_price == actual_ema_price, 0);
+            assert!(price_identifier::get_bytes(&price_identifier) == price_identifier::get_bytes(&actual_price_identifier), 0);
+
+            i = i + 1;
+        };
+    }
+
+    #[test]
+    fun test_get_update_fee() {
+        let single_update_fee = 50;
+        let (scenario, test_coins) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", vector[], single_update_fee, 0);
+        test_scenario::next_tx(&mut scenario, DEPLOYER, );
+        let pyth_state = take_shared<PythState>(&scenario);
+        // Pass in a single VAA
+        assert!(pyth::get_total_update_fee(&pyth_state, &vector[
+            x"fb1543888001083cf2e6ef3afdcf827e89b11efd87c563638df6e1995ada9f93",
+        ]) == single_update_fee, 1);
+
+        // Pass in multiple VAAs
+        assert!(pyth::get_total_update_fee(&pyth_state, &vector[
+            x"4ee17a1a4524118de513fddcf82b77454e51be5d6fc9e29fc72dd6c204c0e4fa",
+            x"c72fdf81cfc939d4286c93fbaaae2eec7bae28a5926fa68646b43a279846ccc1",
+            x"d9a8123a793529c31200339820a3210059ecace6c044f81ecad62936e47ca049",
+            x"84e4f21b3e65cef47fda25d15b4eddda1edf720a1d062ccbf441d6396465fbe6",
+            x"9e73f9041476a93701a0b9c7501422cc2aa55d16100bec628cf53e0281b6f72f"
+        ]) == 250, 1);
+
+        return_shared(pyth_state);
+        coin::burn_for_testing<SUI>(test_coins);
+        test_scenario::end(scenario);
+    }
+
+    #[test]
+    #[expected_failure(abort_code = wormhole::vaa::E_WRONG_VERSION)]
+    fun test_create_price_feeds_corrupt_vaa() {
+        let (scenario, test_coins) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", vector[], 50, 0);
+        test_scenario::next_tx(&mut scenario, DEPLOYER);
+        let pyth_state = take_shared<PythState>(&scenario);
+        let worm_state = take_shared<WormState>(&scenario);
+        let clock = take_shared<Clock>(&scenario);
+
+        // Pass in a corrupt VAA, which should fail deseriaizing
+        let corrupt_vaa = x"90F8bf6A479f320ead074411a4B0e7944Ea8c9C1";
+
+        // Create Pyth price feed
+        pyth::create_price_feeds(
+            &mut worm_state,
+            &mut pyth_state,
+            vector[corrupt_vaa],
+            &clock,
+            ctx(&mut scenario)
+        );
+
+        return_shared(pyth_state);
+        return_shared(worm_state);
+        return_shared(clock);
+        coin::burn_for_testing<SUI>(test_coins);
+        test_scenario::end(scenario);
+    }
+
+    #[test]
+    #[expected_failure(abort_code = pyth::pyth::E_INVALID_DATA_SOURCE)]
+    fun test_create_price_feeds_invalid_data_source() {
+        // Initialize the contract with some valid data sources, excluding our test VAA's source
+        let data_sources = vector<DataSource>[
+            data_source::new(
+                4, external_address::new(bytes32::new(x"0000000000000000000000000000000000000000000000000000000000007742"))
+            ),
+            data_source::new(
+                5, external_address::new(bytes32::new(x"0000000000000000000000000000000000000000000000000000000000007637"))
+            )
+        ];
+
+        let (scenario, test_coins) = setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, 50, 0);
+        test_scenario::next_tx(&mut scenario, DEPLOYER);
+
+        let pyth_state = take_shared<PythState>(&scenario);
+        let worm_state = take_shared<WormState>(&scenario);
+        let clock = take_shared<Clock>(&scenario);
+
+        pyth::create_price_feeds(
+            &mut worm_state,
+            &mut pyth_state,
+            TEST_VAAS,
+            &clock,
+            ctx(&mut scenario)
+        );
+
+        return_shared(pyth_state);
+        return_shared(worm_state);
+        return_shared(clock);
+        coin::burn_for_testing<SUI>(test_coins);
+        test_scenario::end(scenario);
+    }
+
+    #[test_only]
+    fun data_sources_for_test_vaa(): vector<DataSource> {
+        // Set some valid data sources, including our test VAA's source
+        vector<DataSource>[
+            data_source::new(
+                1, external_address::new(bytes32::from_bytes(x"0000000000000000000000000000000000000000000000000000000000000004"))),
+                data_source::new(
+                5, external_address::new(bytes32::new(x"0000000000000000000000000000000000000000000000000000000000007637"))),
+                data_source::new(
+                17, external_address::new(bytes32::new(x"71f8dcb863d176e2c420ad6610cf687359612b6fb392e0642b0ca6b1f186aa3b")))
+        ]
+    }
+
+    #[test]
+    fun test_create_and_update_price_feeds_success() {
+        let data_sources = data_sources_for_test_vaa();
+        let base_update_fee = 50;
+        let coins_to_mint = 5000;
+
+        let (scenario, test_coins) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
+        test_scenario::next_tx(&mut scenario, DEPLOYER);
+
+        let pyth_state = take_shared<PythState>(&scenario);
+        let worm_state = take_shared<WormState>(&scenario);
+        let clock = take_shared<Clock>(&scenario);
+
+        pyth::create_price_feeds(
+            &mut worm_state,
+            &mut pyth_state,
+            TEST_VAAS,
+            &clock,
+            ctx(&mut scenario)
+        );
+
+        // Affirm that 4 objects, which correspond to the 4 new price info objects
+        // containing the price feeds were created and shared.
+        let effects = test_scenario::next_tx(&mut scenario, DEPLOYER);
+        let shared_ids = test_scenario::shared(&effects);
+        let created_ids = test_scenario::created(&effects);
+        assert!(vector::length<ID>(&shared_ids)==4, 0);
+        assert!(vector::length<ID>(&created_ids)==4, 0);
+
+        let price_info_object_1 = take_shared<PriceInfoObject>(&scenario);
+        let price_info_object_2 = take_shared<PriceInfoObject>(&scenario);
+        let price_info_object_3 = take_shared<PriceInfoObject>(&scenario);
+        let price_info_object_4 = take_shared<PriceInfoObject>(&scenario);
+
+        // Create vector of price info objects (Sui objects with key ability and living in global store),
+        // which contain the price feeds we want to update. Note that these can be passed into
+        // update_price_feeds in any order!
+        let price_info_object_vec = vector[price_info_object_1, price_info_object_2, price_info_object_3, price_info_object_4];
+
+        pyth::update_price_feeds(
+            &mut worm_state,
+            &mut pyth_state,
+            TEST_VAAS,
+            &mut price_info_object_vec,
+            test_coins,
+            &clock
+        );
+
+        price_info_object_4 = vector::pop_back(&mut price_info_object_vec);
+        price_info_object_3 = vector::pop_back(&mut price_info_object_vec);
+        price_info_object_2 = vector::pop_back(&mut price_info_object_vec);
+        price_info_object_1 = vector::pop_back(&mut price_info_object_vec);
+        vector::destroy_empty(price_info_object_vec);
+
+        return_shared(pyth_state);
+        return_shared(worm_state);
+        return_shared(price_info_object_1);
+        return_shared(price_info_object_2);
+        return_shared(price_info_object_3);
+        return_shared(price_info_object_4);
+
+        return_shared(clock);
+        test_scenario::end(scenario);
+    }
+
+    #[test]
+    #[expected_failure(abort_code = pyth::pyth::E_PRICE_INFO_OBJECT_NOT_FOUND)]
+    fun test_create_and_update_price_feeds_price_info_object_not_found_failure() {
+        let data_sources = data_sources_for_test_vaa();
+        let base_update_fee = 50;
+        let coins_to_mint = 5000;
+
+        let (scenario, test_coins) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
+        test_scenario::next_tx(&mut scenario, DEPLOYER);
+
+        let pyth_state = take_shared<PythState>(&scenario);
+        let worm_state = take_shared<WormState>(&scenario);
+        let clock = take_shared<Clock>(&scenario);
+
+        pyth::create_price_feeds(
+            &mut worm_state,
+            &mut pyth_state,
+            TEST_VAAS,
+            &clock,
+            ctx(&mut scenario)
+        );
+
+        // Affirm that 4 objects, which correspond to the 4 new price info objects
+        // containing the price feeds were created and shared.
+        let effects = test_scenario::next_tx(&mut scenario, DEPLOYER);
+        let shared_ids = test_scenario::shared(&effects);
+        let created_ids = test_scenario::created(&effects);
+        assert!(vector::length<ID>(&shared_ids)==4, 0);
+        assert!(vector::length<ID>(&created_ids)==4, 0);
+
+        let price_info_object_1 = take_shared<PriceInfoObject>(&scenario);
+        let price_info_object_2 = take_shared<PriceInfoObject>(&scenario);
+        let price_info_object_3 = take_shared<PriceInfoObject>(&scenario);
+        let price_info_object_4 = take_shared<PriceInfoObject>(&scenario);
+
+        // Note that here we only pass in 3 price info objects corresponding to 3 out
+        // of the 4 price feeds.
+        let price_info_object_vec = vector[price_info_object_1, price_info_object_2, price_info_object_3];
+
+        pyth::update_price_feeds(
+            &mut worm_state,
+            &mut pyth_state,
+            TEST_VAAS,
+            &mut price_info_object_vec,
+            test_coins,
+            &clock
+        );
+
+        price_info_object_3 = vector::pop_back(&mut price_info_object_vec);
+        price_info_object_2 = vector::pop_back(&mut price_info_object_vec);
+        price_info_object_1 = vector::pop_back(&mut price_info_object_vec);
+
+        vector::destroy_empty(price_info_object_vec);
+        return_shared(pyth_state);
+        return_shared(worm_state);
+        return_shared(price_info_object_1);
+        return_shared(price_info_object_2);
+        return_shared(price_info_object_3);
+        return_shared(price_info_object_4);
+
+        return_shared(clock);
+        test_scenario::end(scenario);
+    }
+
+    #[test]
+    #[expected_failure(abort_code = pyth::pyth::E_INSUFFICIENT_FEE)]
+    fun test_create_and_update_price_feeds_insufficient_fee() {
+        let data_sources = data_sources_for_test_vaa();
+        let base_update_fee = 50;
+        let coins_to_mint = 5;
+
+        let (scenario, test_coins) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
+        test_scenario::next_tx(&mut scenario, DEPLOYER);
+
+        let pyth_state = take_shared<PythState>(&scenario);
+        let worm_state = take_shared<WormState>(&scenario);
+        let clock = take_shared<Clock>(&scenario);
+
+        pyth::create_price_feeds(
+            &mut worm_state,
+            &mut pyth_state,
+            TEST_VAAS,
+            &clock,
+            ctx(&mut scenario)
+        );
+
+        test_scenario::next_tx(&mut scenario, DEPLOYER);
+
+        let price_info_object = take_shared<PriceInfoObject>(&scenario);
+        let price_info_object_vec = vector[price_info_object];
+
+        pyth::update_price_feeds(
+            &mut worm_state,
+            &mut pyth_state,
+            TEST_VAAS,
+            &mut price_info_object_vec,
+            test_coins,
+            &clock
+        );
+
+        price_info_object = vector::pop_back(&mut price_info_object_vec);
+        vector::destroy_empty(price_info_object_vec);
+        return_shared(pyth_state);
+        return_shared(worm_state);
+        return_shared(price_info_object);
+        return_shared(clock);
+        test_scenario::end(scenario);
+    }
+
+    #[test]
+    fun test_update_cache(){
+        let data_sources = data_sources_for_test_vaa();
+        let base_update_fee = 50;
+        let coins_to_mint = 5000;
+
+        let (scenario, test_coins) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
+        test_scenario::next_tx(&mut scenario, DEPLOYER);
+
+        let pyth_state = take_shared<PythState>(&scenario);
+        let worm_state = take_shared<WormState>(&scenario);
+        let clock = take_shared<Clock>(&scenario);
+
+        pyth::create_price_feeds(
+            &mut worm_state,
+            &mut pyth_state,
+            TEST_VAAS,
+            &clock,
+            ctx(&mut scenario)
+        );
+
+        test_scenario::next_tx(&mut scenario, DEPLOYER);
+
+        let price_info_object_1 = take_shared<PriceInfoObject>(&scenario);
+        let price_info_object_2 = take_shared<PriceInfoObject>(&scenario);
+        let price_info_object_3 = take_shared<PriceInfoObject>(&scenario);
+        let price_info_object_4 = take_shared<PriceInfoObject>(&scenario);
+
+        // These updates are price infos that correspond to the ones in TEST_VAAS.
+        let updates = get_mock_price_infos();
+        let price_info_object_vec = vector[
+            price_info_object_1,
+            price_info_object_2,
+            price_info_object_3,
+            price_info_object_4
+        ];
+
+        // Check that TEST_VAAS was indeed used to instantiate the price feeds correctly,
+        // by confirming that the info in updates is contained in price_info_object_vec.
+        check_price_feeds_cached(&updates, &price_info_object_vec);
+
+        price_info_object_4 = vector::pop_back(&mut price_info_object_vec);
+        price_info_object_3 = vector::pop_back(&mut price_info_object_vec);
+        price_info_object_2 = vector::pop_back(&mut price_info_object_vec);
+        price_info_object_1 = vector::pop_back(&mut price_info_object_vec);
+        vector::destroy_empty(price_info_object_vec);
+
+        return_shared(pyth_state);
+        return_shared(worm_state);
+        return_shared(price_info_object_1);
+        return_shared(price_info_object_2);
+        return_shared(price_info_object_3);
+        return_shared(price_info_object_4);
+        coin::burn_for_testing<SUI>(test_coins);
+
+        return_shared(clock);
+        test_scenario::end(scenario);
+    }
+
+    #[test]
+    fun test_update_cache_old_update() {
+        let data_sources = data_sources_for_test_vaa();
+        let base_update_fee = 50;
+        let coins_to_mint = 5000;
+
+        let (scenario, test_coins) =  setup_test(500, 23, x"5d1f252d5de865279b00c84bce362774c2804294ed53299bc4a0389a5defef92", data_sources, base_update_fee, coins_to_mint);
+        test_scenario::next_tx(&mut scenario, DEPLOYER);
+
+        let pyth_state = take_shared<PythState>(&scenario);
+        let worm_state = take_shared<WormState>(&scenario);
+        let clock = take_shared<Clock>(&scenario);
+
+        pyth::create_price_feeds(
+            &mut worm_state,
+            &mut pyth_state,
+            TEST_VAAS,
+            &clock,
+            ctx(&mut scenario)
+        );
+
+        test_scenario::next_tx(&mut scenario, DEPLOYER);
+
+        let price_info_object_1 = take_shared<PriceInfoObject>(&scenario);
+        let price_info_object_2 = take_shared<PriceInfoObject>(&scenario);
+        let price_info_object_3 = take_shared<PriceInfoObject>(&scenario);
+        let price_info_object_4 = take_shared<PriceInfoObject>(&scenario);
+
+        let price_info_object_vec = vector[
+            price_info_object_1,
+            price_info_object_2,
+            price_info_object_3,
+            price_info_object_4
+        ];
+
+        // Hardcode the price identifier, price, and ema_price for price_info_object_1, because
+        // it's easier than unwrapping price_info_object_1 and getting the quantities via getters.
+        let timestamp = 1663680740;
+        let price_identifier = price_identifier::from_byte_vec(x"c6c75c89f14810ec1c54c03ab8f1864a4c4032791f05747f560faec380a695d1");
+        let price = price::new(i64::new(1557, false), 7, i64::new(5, true), timestamp);
+        let ema_price = price::new(i64::new(1500, false), 3, i64::new(5, true), timestamp);
+
+        // Attempt to update the price with an update older than the current cached one.
+        let old_price = price::new(i64::new(1243, true), 9802, i64::new(6, false), timestamp - 200);
+        let old_ema_price = price::new(i64::new(8976, true), 234, i64::new(897, false), timestamp - 200);
+        let old_update = price_info::new_price_info(
+            1257278600,
+            1690226180,
+            price_feed::new(
+                    price_identifier,
+                    old_price,
+                    old_ema_price,
+            )
+        );
+        pyth::update_cache(vector<PriceInfo>[old_update], &mut price_info_object_vec, &clock);
+
+        price_info_object_4 = vector::pop_back(&mut price_info_object_vec);
+        price_info_object_3 = vector::pop_back(&mut price_info_object_vec);
+        price_info_object_2 = vector::pop_back(&mut price_info_object_vec);
+        price_info_object_1 = vector::pop_back(&mut price_info_object_vec);
+
+        vector::destroy_empty(price_info_object_vec);
+
+        // Confirm that the current price and ema price didn't change
+        let current_price_info = price_info::get_price_info_from_price_info_object(&price_info_object_1);
+        let current_price_feed = price_info::get_price_feed(&current_price_info);
+        let current_price = price_feed::get_price(current_price_feed);
+        let current_ema_price = price_feed::get_ema_price(current_price_feed);
+
+        // Confirm that no price update occurred when we tried to update cache with an
+        // outdated update: old_update.
+        assert!(current_price == price, 1);
+        assert!(current_ema_price == ema_price, 1);
+
+        test_scenario::next_tx(&mut scenario, DEPLOYER);
+
+        // Update the cache with a fresh update.
+        let fresh_price = price::new(i64::new(5243, true), 2, i64::new(3, false), timestamp + 200);
+        let fresh_ema_price = price::new(i64::new(8976, true), 21, i64::new(32, false), timestamp + 200);
+        let fresh_update = price_info::new_price_info(
+            1257278600,
+            1690226180,
+            price_feed::new(
+                    price_identifier,
+                    fresh_price,
+                    fresh_ema_price,
+            )
+        );
+
+        let price_info_object_vec = vector[
+            price_info_object_1,
+            price_info_object_2,
+            price_info_object_3,
+            price_info_object_4
+        ];
+
+        pyth::update_cache(vector<PriceInfo>[fresh_update], &mut price_info_object_vec, &clock);
+
+        price_info_object_4 = vector::pop_back(&mut price_info_object_vec);
+        price_info_object_3 = vector::pop_back(&mut price_info_object_vec);
+        price_info_object_2 = vector::pop_back(&mut price_info_object_vec);
+        price_info_object_1 = vector::pop_back(&mut price_info_object_vec);
+
+        vector::destroy_empty(price_info_object_vec);
+
+        // Confirm that the Pyth cached price got updated to fresh_price
+        let current_price_info = price_info::get_price_info_from_price_info_object(&price_info_object_1);
+        let current_price_feed = price_info::get_price_feed(&current_price_info);
+        let current_price = price_feed::get_price(current_price_feed);
+        let current_ema_price = price_feed::get_ema_price(current_price_feed);
+
+        assert!(current_price==fresh_price, 0);
+        assert!(current_ema_price==fresh_ema_price, 0);
+
+        return_shared(pyth_state);
+        return_shared(worm_state);
+        return_shared(price_info_object_1);
+        return_shared(price_info_object_2);
+        return_shared(price_info_object_3);
+        return_shared(price_info_object_4);
+        coin::burn_for_testing<SUI>(test_coins);
+
+        return_shared(clock);
+        test_scenario::end(scenario);
+    }
+}
+
 // TODO - pyth tests
 // https://github.com/pyth-network/pyth-crosschain/blob/main/target_chains/aptos/contracts/sources/pyth.move#L384

target_chains/sui/contracts/sources/state.move

@@ -3,12 +3,16 @@ module pyth::state {
     use sui::object::{Self, UID, ID};
     use sui::transfer::{Self};
     use sui::tx_context::{Self, TxContext};
+    use sui::package::{Self, UpgradeCap};
 
     use pyth::data_source::{Self, DataSource};
     use pyth::price_info::{Self};
     use pyth::price_identifier::{PriceIdentifier};
 
+    use wormhole::setup::{assert_package_upgrade_cap};
+
     friend pyth::pyth;
+    friend pyth::pyth_tests;
     friend pyth::governance_action;
     friend pyth::set_update_fee;
     friend pyth::set_stale_price_threshold;
@@ -28,6 +32,7 @@ module pyth::state {
         last_executed_governance_sequence: u64,
         stale_price_threshold: u64,
         base_update_fee: u64,
+        upgrade_cap: UpgradeCap
     }
 
     fun init(ctx: &mut TxContext) {
@@ -39,18 +44,41 @@ module pyth::state {
         );
     }
 
+    #[test_only]
+    public fun init_test_only(ctx: &mut TxContext) {
+        init(ctx);
+
+        // This will be created and sent to the transaction sender
+        // automatically when the contract is published.
+        transfer::public_transfer(
+            sui::package::test_publish(object::id_from_address(@pyth), ctx),
+            tx_context::sender(ctx)
+        );
+    }
+
     // Initialization
     public(friend) fun init_and_share_state(
         deployer: DeployerCap,
+        upgrade_cap: UpgradeCap,
         stale_price_threshold: u64,
         base_update_fee: u64,
         governance_data_source: DataSource,
         sources: vector<DataSource>,
         ctx: &mut TxContext
     ) {
+        // TODO - version control
+        // let version = wormhole::version_control::version();
+        //assert!(version == 1, E_INVALID_BUILD_VERSION);
+
         let DeployerCap { id } = deployer;
         object::delete(id);
 
+        assert_package_upgrade_cap<DeployerCap>(
+            &upgrade_cap,
+            package::compatible_policy(),
+            1 // version
+        );
+
         let uid = object::new(ctx);
 
         // Create a set that contains all registered data sources and
@@ -73,6 +101,7 @@ module pyth::state {
         transfer::share_object(
             State {
                 id: uid,
+                upgrade_cap,
                 governance_data_source,
                 last_executed_governance_sequence: 0,
                 stale_price_threshold,