feat(target_chains/ton): add execute_governance_action (#1882)

* add execute_governance_action

* add more governance test payload

* fix comment for ACCUMULATOR_WORMHOLE_MAGIC

* remove other constants

* add test for set data sources governance instruction

* refactor

* add execute_authorize_governance_data_source_transfer

* precommit

* make set_valid_period unsupported

* add unhappy path for update price feeds with insufficient fee

* precommit

* add more unhappy paths

* add more tests

* add comment

target_chains/ton/contracts/contracts/Main.fc

@@ -1,13 +1,10 @@
 #include "imports/stdlib.fc";
 #include "common/errors.fc";
 #include "common/storage.fc";
+#include "common/op.fc";
 #include "Wormhole.fc";
 #include "Pyth.fc";
 
-;; Opcodes
-const int OP_UPDATE_GUARDIAN_SET = 1;
-const int OP_EXECUTE_GOVERNANCE_ACTION = 2;
-
 ;; Internal message handler
 () recv_internal(int my_balance, int msg_value, cell in_msg_full, slice in_msg_body) impure {
     if (in_msg_body.slice_empty?()) { ;; ignore empty messages

target_chains/ton/contracts/contracts/Pyth.fc

@@ -4,6 +4,8 @@
 #include "common/utils.fc";
 #include "common/constants.fc";
 #include "common/merkle_tree.fc";
+#include "common/governance_actions.fc";
+#include "common/gas.fc";
 #include "./Wormhole.fc";
 
 cell store_price(int price, int conf, int expo, int publish_time) {
@@ -59,7 +61,7 @@ slice read_and_verify_header(slice data) {
     return (price, conf, expo, publish_time);
 }
 
-(int) get_update_fee(slice data) method_id {
+int get_update_fee(slice data) method_id {
     load_data();
     slice cs = read_and_verify_header(data);
     int wormhole_proof_size_bytes = cs~load_uint(16);
@@ -68,6 +70,11 @@ slice read_and_verify_header(slice data) {
     return single_update_fee * num_updates;
 }
 
+int get_single_update_fee() method_id {
+    load_data();
+    return single_update_fee;
+}
+
 int get_governance_data_source_index() method_id {
     load_data();
     return governance_data_source_index;
@@ -83,6 +90,17 @@ int get_last_executed_governance_sequence() method_id {
     return last_executed_governance_sequence;
 }
 
+int get_is_valid_data_source(cell data_source) method_id {
+    load_data();
+    int data_source_key = cell_hash(data_source);
+    (slice value, int found?) = is_valid_data_source.udict_get?(256, data_source_key);
+    if (found?) {
+        return value~load_int(1);
+    } else {
+        return 0;
+    }
+}
+
 (int, int, int, int) get_price_unsafe(int price_feed_id) method_id {
     load_data();
     (slice result, int success) = latest_price_feeds.udict_get?(256, price_feed_id);
@@ -148,10 +166,13 @@ int parse_pyth_payload_in_wormhole_vm(slice payload) impure {
     cs = new_cs;
 
     int num_updates = cs~load_uint(8);
-    int fee = single_update_fee * num_updates;
+    int update_fee = single_update_fee * num_updates;
+    int compute_fee = get_compute_fee(WORKCHAIN, UPDATE_PRICE_FEEDS_GAS);
+    throw_unless(ERROR_INSUFFICIENT_GAS, msg_value >= compute_fee);
+    int remaining_msg_value = msg_value - compute_fee;
 
-    ;; Check if the sender has sent enough TON to cover the fee
-    throw_unless(ERROR_INSUFFICIENT_FEE, msg_value >= fee);
+    ;; Check if the sender has sent enough TON to cover the update_fee
+    throw_unless(ERROR_INSUFFICIENT_FEE, remaining_msg_value >= update_fee);
 
     (_, _, _, _, int emitter_chain_id, int emitter_address, _, _, slice payload, _) = parse_and_verify_wormhole_vm(wormhole_proof.begin_parse());
 
@@ -203,6 +224,131 @@ int parse_pyth_payload_in_wormhole_vm(slice payload) impure {
     store_data();
 }
 
-() execute_governance_action(slice in_msg_body) impure {
+() verify_governance_vm(int emitter_chain_id, int emitter_address, int sequence) impure {
+    (int gov_chain_id, int gov_address) = parse_data_source(governance_data_source);
+    throw_unless(ERROR_INVALID_GOVERNANCE_DATA_SOURCE, emitter_chain_id == gov_chain_id);
+    throw_unless(ERROR_INVALID_GOVERNANCE_DATA_SOURCE, emitter_address == gov_address);
+    throw_if(ERROR_OLD_GOVERNANCE_MESSAGE, sequence <= last_executed_governance_sequence);
+    last_executed_governance_sequence = sequence;
+}
+
+(int, int, slice) parse_governance_instruction(slice payload) {
+    int magic = payload~load_uint(32);
+    throw_unless(ERROR_INVALID_GOVERNANCE_MAGIC, magic == GOVERNANCE_MAGIC);
+
+    int module = payload~load_uint(8);
+    throw_unless(ERROR_INVALID_GOVERNANCE_MODULE, module == GOVERNANCE_MODULE);
+
+    int action = payload~load_uint(8);
+
+    int target_chain_id = payload~load_uint(16);
+
+    return (target_chain_id, action, payload);
+}
+
+int apply_decimal_expo(int value, int expo) {
+    int result = value;
+    repeat (expo) {
+        result *= 10;
+    }
+    return result;
+}
+
+() execute_upgrade_contract(slice payload) impure {
     ;; TODO: Implement
 }
+
+() execute_authorize_governance_data_source_transfer(slice payload) impure {
+    ;; Verify the claim VAA
+    (_, _, _, _, int emitter_chain_id, int emitter_address, int sequence, _, slice claim_payload, _) = parse_and_verify_wormhole_vm(payload);
+
+    ;; Parse the claim payload
+    (int target_chain_id, int action, slice claim_payload) = parse_governance_instruction(claim_payload);
+
+    ;; Verify that this is a valid governance action for this chain
+    throw_if(ERROR_INVALID_GOVERNANCE_TARGET, (target_chain_id != 0) & (target_chain_id != chain_id));
+    throw_unless(ERROR_INVALID_GOVERNANCE_ACTION, action == REQUEST_GOVERNANCE_DATA_SOURCE_TRANSFER);
+
+    ;; Extract the new governance data source index from the claim payload
+    int new_governance_data_source_index = claim_payload~load_uint(32);
+
+    ;; Verify that the new index is greater than the current index
+    int current_index = governance_data_source_index;
+    throw_if(ERROR_OLD_GOVERNANCE_MESSAGE, new_governance_data_source_index <= current_index);
+
+    ;; Update the governance data source
+    governance_data_source = begin_cell()
+        .store_uint(emitter_chain_id, 16)
+        .store_uint(emitter_address, 256)
+    .end_cell();
+
+    governance_data_source_index = new_governance_data_source_index;
+
+    ;; Update the last executed governance sequence
+    last_executed_governance_sequence = sequence;
+}
+
+() execute_set_data_sources(slice payload) impure {
+    int num_sources = payload~load_uint(8);
+    cell new_data_sources = new_dict();
+
+    repeat(num_sources) {
+        (cell data_source, slice new_payload) = read_and_store_large_data(payload, 272); ;; 272 = 256 + 16
+        payload = new_payload;
+        slice data_source_slice = data_source.begin_parse();
+        int emitter_chain_id = data_source_slice~load_uint(16);
+        int emitter_address = data_source_slice~load_uint(256);
+        cell data_source = begin_cell()
+            .store_uint(emitter_chain_id, 16)
+            .store_uint(emitter_address, 256)
+        .end_cell();
+        int data_source_key = cell_hash(data_source);
+        new_data_sources~udict_set(256, data_source_key, begin_cell().store_int(true, 1).end_cell().begin_parse());
+    }
+
+    is_valid_data_source = new_data_sources;
+}
+
+() execute_set_fee(slice payload) impure {
+    int value = payload~load_uint(64);
+    int expo = payload~load_uint(64);
+    int new_fee = apply_decimal_expo(value, expo);
+    single_update_fee = new_fee;
+}
+
+() execute_governance_payload(int action, slice payload) impure {
+    if (action == UPGRADE_CONTRACT) {
+        execute_upgrade_contract(payload);
+    } elseif (action == AUTHORIZE_GOVERNANCE_DATA_SOURCE_TRANSFER) {
+        execute_authorize_governance_data_source_transfer(payload);
+    } elseif (action == SET_DATA_SOURCES) {
+        execute_set_data_sources(payload);
+    } elseif (action == SET_FEE) {
+        execute_set_fee(payload);
+    } elseif (action == SET_VALID_PERIOD) {
+        ;; Unsupported governance action
+        throw(ERROR_INVALID_GOVERNANCE_ACTION);
+    } elseif (action == REQUEST_GOVERNANCE_DATA_SOURCE_TRANSFER) {
+        ;; RequestGovernanceDataSourceTransfer can only be part of
+        ;; AuthorizeGovernanceDataSourceTransfer message
+        throw(ERROR_INVALID_GOVERNANCE_ACTION);
+    } else {
+        throw(ERROR_INVALID_GOVERNANCE_ACTION);
+    }
+}
+
+() execute_governance_action(slice in_msg_body) impure {
+    load_data();
+
+    (_, _, _, _, int emitter_chain_id, int emitter_address, int sequence, _, slice payload, _) = parse_and_verify_wormhole_vm(in_msg_body);
+
+    verify_governance_vm(emitter_chain_id, emitter_address, sequence);
+
+    (int target_chain_id, int action, slice payload) = parse_governance_instruction(payload);
+
+    throw_if(ERROR_INVALID_GOVERNANCE_TARGET, (target_chain_id != 0) & (target_chain_id != chain_id));
+
+    execute_governance_payload(action, payload);
+
+    store_data();
+}

target_chains/ton/contracts/contracts/common/constants.fc

@@ -1,5 +1,7 @@
 const int ACCUMULATOR_MAGIC = 0x504e4155; ;; "PNAU" (Pyth Network Accumulator Update)
-const int ACCUMULATOR_WORMHOLE_MAGIC = 0x41555756; ;; Stands for AUWV (Accumulator Update Wormhole Verficiation)
+const int ACCUMULATOR_WORMHOLE_MAGIC = 0x41555756; ;; "AUWV" (Accumulator Update Wormhole Verficiation)
+const int GOVERNANCE_MAGIC = 0x5054474d; ;; "PTGM" (Pyth Governance Message)
+const int GOVERNANCE_MODULE = 1;
 const int MAJOR_VERSION = 1;
 const int MINIMUM_ALLOWED_MINOR_VERSION = 0;
 
@@ -7,3 +9,17 @@ const int GUARDIAN_SET_EXPIRY = 86400; ;; 1 day in seconds
 const int UPGRADE_MODULE = 0x0000000000000000000000000000000000000000000000000000000000436f7265; ;; "Core" (left-padded to 256 bits) in hex
 
 const int WORMHOLE_MERKLE_UPDATE_TYPE = 0;
+
+{-
+  The main workchain ID in TON. Currently, TON has two blockchains:
+  1. Masterchain: Used for system-level operations and consensus.
+  2. Basechain/Workchain: The primary chain for user accounts and smart contracts.
+
+  While TON supports up to 2^32 workchains, currently only Workchain 0 is active.
+  This constant defines the default workchain for smart contract deployment and interactions.
+
+  Note: Gas costs differ between chains:
+  - Basechain: 1 gas = 400 nanotons = 0.0000004 TON
+  - Masterchain: 1 gas = 10000 nanotons = 0.00001 TON (25x more expensive)
+-}
+const int WORKCHAIN = 0;

target_chains/ton/contracts/contracts/common/errors.fc

@@ -35,3 +35,11 @@ const int ERROR_INVALID_UPDATE_DATA_TYPE = 1028;
 const int ERROR_INVALID_MESSAGE_TYPE = 1029;
 const int ERROR_INSUFFICIENT_FEE = 1030;
 const int ERROR_INVALID_PROOF_SIZE = 1031;
+const int ERROR_INVALID_GOVERNANCE_DATA_SOURCE = 1032;
+const int ERROR_OLD_GOVERNANCE_MESSAGE = 1033;
+const int ERROR_INVALID_GOVERNANCE_TARGET = 1034;
+const int ERROR_INVALID_GOVERNANCE_MAGIC = 1035;
+const int ERROR_INVALID_GOVERNANCE_MODULE = 1036;
+
+;; Common
+const int ERROR_INSUFFICIENT_GAS = 1037;

target_chains/ton/contracts/contracts/common/gas.fc

@@ -0,0 +1,4 @@
+int get_compute_fee(int workchain, int gas_used) asm(gas_used workchain) "GETGASFEE";
+
+;; The actual gas used for the transaction is 350166 but we add ~10% (385182.6) and round up (390000) to be on the safe side because the amount of gas used can vary based on the current state of the blockchain
+const int UPDATE_PRICE_FEEDS_GAS = 390000;

target_chains/ton/contracts/contracts/common/governance_actions.fc

@@ -0,0 +1,6 @@
+const int UPGRADE_CONTRACT = 0;
+const int AUTHORIZE_GOVERNANCE_DATA_SOURCE_TRANSFER = 1;
+const int SET_DATA_SOURCES = 2;
+const int SET_FEE = 3;
+const int SET_VALID_PERIOD = 4;
+const int REQUEST_GOVERNANCE_DATA_SOURCE_TRANSFER = 5;

target_chains/ton/contracts/contracts/common/op.fc

@@ -1,2 +1,3 @@
 const int OP_UPDATE_GUARDIAN_SET = 1;
 const int OP_UPDATE_PRICE_FEEDS = 2;
+const int OP_EXECUTE_GOVERNANCE_ACTION = 3;

target_chains/ton/contracts/contracts/tests/PythTest.fc

@@ -24,6 +24,8 @@
         update_guardian_set(data.begin_parse());
     } elseif (op == OP_UPDATE_PRICE_FEEDS) {
         update_price_feeds(msg_value, data.begin_parse());
+    } elseif (op == OP_EXECUTE_GOVERNANCE_ACTION) {
+        execute_governance_action(data.begin_parse());
     } else {
         throw(0xffff); ;; Throw exception for unknown op
     }
@@ -49,6 +51,10 @@
     return get_update_fee(in_msg_body);
 }
 
+(int) test_get_single_update_fee() method_id {
+    return get_single_update_fee();
+}
+
 (int) test_get_chain_id() method_id {
     return get_chain_id();
 }
@@ -64,3 +70,7 @@
 (cell) test_get_governance_data_source() method_id {
     return get_governance_data_source();
 }
+
+(int) test_get_is_valid_data_source(cell data_source) method_id {
+    return get_is_valid_data_source(data_source);
+}

target_chains/ton/contracts/tests/PythTest.spec.ts

@@ -4,9 +4,18 @@ import "@ton/test-utils";
 import { compile } from "@ton/blueprint";
 import { HexString, Price } from "@pythnetwork/price-service-sdk";
 import { PythTest, PythTestConfig } from "../wrappers/PythTest";
-import { BTC_PRICE_FEED_ID, HERMES_BTC_ETH_UPDATE } from "./utils/pyth";
+import {
+  BTC_PRICE_FEED_ID,
+  HERMES_BTC_ETH_UPDATE,
+  PYTH_SET_DATA_SOURCES,
+  PYTH_SET_FEE,
+  TEST_GUARDIAN_ADDRESS1,
+  PYTH_AUTHORIZE_GOVERNANCE_DATA_SOURCE_TRANSFER,
+  PYTH_REQUEST_GOVERNANCE_DATA_SOURCE_TRANSFER,
+} from "./utils/pyth";
 import { GUARDIAN_SET_0, MAINNET_UPGRADE_VAAS } from "./utils/wormhole";
 import { DataSource } from "@pythnetwork/xc-admin-common";
+import { parseDataSource } from "./utils";
 
 const TIME_PERIOD = 60;
 const PRICE = new Price({
@@ -29,6 +38,18 @@ const DATA_SOURCES: DataSource[] = [
       "e101faedac5851e32b9b23b5f9411a8c2bac4aae3ed4dd7b811dd1a72ea4aa71",
   },
 ];
+const TEST_GOVERNANCE_DATA_SOURCES: DataSource[] = [
+  {
+    emitterChain: 1,
+    emitterAddress:
+      "0000000000000000000000000000000000000000000000000000000000000029",
+  },
+  {
+    emitterChain: 2,
+    emitterAddress:
+      "000000000000000000000000000000000000000000000000000000000000002b",
+  },
+];
 
 describe("PythTest", () => {
   let code: Cell;
@@ -57,7 +78,8 @@ describe("PythTest", () => {
     guardianSet: string[] = GUARDIAN_SET_0,
     chainId: number = 1,
     governanceChainId: number = 1,
-    governanceContract: string = "0000000000000000000000000000000000000000000000000000000000000004"
+    governanceContract: string = "0000000000000000000000000000000000000000000000000000000000000004",
+    governanceDataSource?: DataSource
   ) {
     const config: PythTestConfig = {
       priceFeedId,
@@ -71,6 +93,7 @@ describe("PythTest", () => {
       chainId,
       governanceChainId,
       governanceContract,
+      governanceDataSource,
     };
 
     pythTest = blockchain.openContract(PythTest.createFromConfig(config, code));
@@ -88,6 +111,23 @@ describe("PythTest", () => {
     });
   }
 
+  async function updateGuardianSets(
+    pythTest: SandboxContract<PythTest>,
+    deployer: SandboxContract<TreasuryContract>
+  ) {
+    for (const vaa of MAINNET_UPGRADE_VAAS) {
+      const result = await pythTest.sendUpdateGuardianSet(
+        deployer.getSender(),
+        Buffer.from(vaa, "hex")
+      );
+      expect(result.transactions).toHaveTransaction({
+        from: deployer.address,
+        to: pythTest.address,
+        success: true,
+      });
+    }
+  }
+
   it("should correctly get price unsafe", async () => {
     await deployContract();
 
@@ -181,105 +221,502 @@ describe("PythTest", () => {
     await deployContract();
     let result;
 
-    const mainnet_upgrade_vaa_1 = MAINNET_UPGRADE_VAAS[0];
-    result = await pythTest.sendUpdateGuardianSet(
+    await updateGuardianSets(pythTest, deployer);
+
+    const updateData = Buffer.from(HERMES_BTC_ETH_UPDATE, "hex");
+    const updateFee = await pythTest.getUpdateFee(updateData);
+
+    result = await pythTest.sendUpdatePriceFeeds(
       deployer.getSender(),
-      Buffer.from(mainnet_upgrade_vaa_1, "hex")
+      updateData,
+      toNano(updateFee)
     );
+
     expect(result.transactions).toHaveTransaction({
       from: deployer.address,
       to: pythTest.address,
       success: true,
     });
 
-    const mainnet_upgrade_vaa_2 = MAINNET_UPGRADE_VAAS[1];
-    result = await pythTest.sendUpdateGuardianSet(
+    // Check if the price has been updated correctly
+    const updatedPrice = await pythTest.getPriceUnsafe(BTC_PRICE_FEED_ID);
+    expect(updatedPrice.price).not.toBe(Number(PRICE.price)); // Since we updated the price, it should not be the same as the initial price
+    expect(updatedPrice.publishTime).toBeGreaterThan(PRICE.publishTime);
+  });
+
+  it("should fail to get update fee with invalid data", async () => {
+    await deployContract();
+    await updateGuardianSets(pythTest, deployer);
+
+    const invalidUpdateData = Buffer.from("invalid data");
+
+    await expect(pythTest.getUpdateFee(invalidUpdateData)).rejects.toThrow(
+      "Unable to execute get method. Got exit_code: 1021"
+    ); // ERROR_INVALID_MAGIC = 1021
+  });
+
+  it("should fail to update price feeds with invalid data", async () => {
+    await deployContract();
+    await updateGuardianSets(pythTest, deployer);
+
+    const invalidUpdateData = Buffer.from("invalid data");
+
+    // Use a fixed value for updateFee since we can't get it from getUpdateFee
+    const updateFee = toNano("0.1"); // Use a reasonable amount
+
+    const result = await pythTest.sendUpdatePriceFeeds(
       deployer.getSender(),
-      Buffer.from(mainnet_upgrade_vaa_2, "hex")
+      invalidUpdateData,
+      updateFee
     );
+
     expect(result.transactions).toHaveTransaction({
       from: deployer.address,
       to: pythTest.address,
-      success: true,
+      success: false,
+      exitCode: 1021, // ERROR_INVALID_MAGIC
     });
+  });
 
-    const mainnet_upgrade_vaa_3 = MAINNET_UPGRADE_VAAS[2];
-    result = await pythTest.sendUpdateGuardianSet(
+  it("should fail to update price feeds with outdated guardian set", async () => {
+    await deployContract();
+    // Don't update guardian sets
+
+    const updateData = Buffer.from(HERMES_BTC_ETH_UPDATE, "hex");
+    const updateFee = await pythTest.getUpdateFee(updateData);
+
+    const result = await pythTest.sendUpdatePriceFeeds(
       deployer.getSender(),
-      Buffer.from(mainnet_upgrade_vaa_3, "hex")
+      updateData,
+      toNano(updateFee)
     );
+
     expect(result.transactions).toHaveTransaction({
       from: deployer.address,
       to: pythTest.address,
-      success: true,
+      success: false,
+      exitCode: 1002, // ERROR_GUARDIAN_SET_NOT_FOUND
     });
+  });
 
-    const mainnet_upgrade_vaa_4 = MAINNET_UPGRADE_VAAS[3];
-    result = await pythTest.sendUpdateGuardianSet(
+  it("should fail to update price feeds with invalid data source", async () => {
+    await deployContract(
+      BTC_PRICE_FEED_ID,
+      TIME_PERIOD,
+      PRICE,
+      EMA_PRICE,
+      SINGLE_UPDATE_FEE,
+      [] // Empty data sources
+    );
+    await updateGuardianSets(pythTest, deployer);
+
+    const updateData = Buffer.from(HERMES_BTC_ETH_UPDATE, "hex");
+    const updateFee = await pythTest.getUpdateFee(updateData);
+
+    const result = await pythTest.sendUpdatePriceFeeds(
       deployer.getSender(),
-      Buffer.from(mainnet_upgrade_vaa_4, "hex")
+      updateData,
+      toNano(updateFee)
     );
+
     expect(result.transactions).toHaveTransaction({
       from: deployer.address,
       to: pythTest.address,
-      success: true,
+      success: false,
+      exitCode: 1024, // ERROR_UPDATE_DATA_SOURCE_NOT_FOUND
+    });
+  });
+
+  it("should fail to update price feeds with insufficient gas", async () => {
+    await deployContract();
+    await updateGuardianSets(pythTest, deployer);
+
+    const updateData = Buffer.from(HERMES_BTC_ETH_UPDATE, "hex");
+
+    const result = await pythTest.sendUpdatePriceFeeds(
+      deployer.getSender(),
+      updateData,
+      toNano("0.1") // Insufficient gas
+    );
+
+    expect(result.transactions).toHaveTransaction({
+      from: deployer.address,
+      to: pythTest.address,
+      success: false,
+      exitCode: 1037, // ERROR_INSUFFICIENT_GAS
     });
+  });
+
+  it("should fail to update price feeds with insufficient fee", async () => {
+    await deployContract();
+
+    await updateGuardianSets(pythTest, deployer);
 
     const updateData = Buffer.from(HERMES_BTC_ETH_UPDATE, "hex");
     const updateFee = await pythTest.getUpdateFee(updateData);
 
-    result = await pythTest.sendUpdatePriceFeeds(
+    // Send less than the required fee
+    const insufficientFee = updateFee - 1;
+
+    const result = await pythTest.sendUpdatePriceFeeds(
       deployer.getSender(),
       updateData,
-      toNano(updateFee)
+      156000000n + BigInt(insufficientFee) // 156000000 = 390000 (estimated gas used for the transaction, this is defined in contracts/common/gas.fc as UPDATE_PRICE_FEEDS_GAS) * 400 (current settings in basechain are as follows: 1 unit of gas costs 400 nanotons)
     );
 
+    // Check that the transaction did not succeed
     expect(result.transactions).toHaveTransaction({
       from: deployer.address,
       to: pythTest.address,
-      success: true,
+      success: false,
+      exitCode: 1030, // ERROR_INSUFFICIENT_FEE = 1030
     });
+  });
 
-    // Check if the price has been updated correctly
-    const updatedPrice = await pythTest.getPriceUnsafe(BTC_PRICE_FEED_ID);
-    expect(updatedPrice.price).not.toBe(Number(PRICE.price)); // Since we updated the price, it should not be the same as the initial price
-    expect(updatedPrice.publishTime).toBeGreaterThan(PRICE.publishTime);
+  it("should fail to get price for non-existent price feed", async () => {
+    await deployContract();
+
+    const nonExistentPriceFeedId =
+      "0000000000000000000000000000000000000000000000000000000000000000";
+
+    await expect(
+      pythTest.getPriceUnsafe(nonExistentPriceFeedId)
+    ).rejects.toThrow("Unable to execute get method. Got exit_code: 1019"); // ERROR_PRICE_FEED_NOT_FOUND = 1019
   });
 
-  it("should return the correct chain ID", async () => {
+  it("should correctly get chain ID", async () => {
     await deployContract();
 
     const result = await pythTest.getChainId();
     expect(result).toEqual(1);
   });
 
-  it("should return the correct last executed governance sequence", async () => {
-    await deployContract();
+  it("should correctly get last executed governance sequence", async () => {
+    await deployContract(
+      BTC_PRICE_FEED_ID,
+      TIME_PERIOD,
+      PRICE,
+      EMA_PRICE,
+      SINGLE_UPDATE_FEE,
+      DATA_SOURCES,
+      0,
+      [TEST_GUARDIAN_ADDRESS1],
+      60051,
+      1,
+      "0000000000000000000000000000000000000000000000000000000000000004",
+      TEST_GOVERNANCE_DATA_SOURCES[0]
+    );
+
+    // Check initial value
+    let result = await pythTest.getLastExecutedGovernanceSequence();
+    expect(result).toEqual(0);
 
-    const result = await pythTest.getLastExecutedGovernanceSequence();
-    expect(result).toEqual(0); // Initial value should be 0
+    // Execute a governance action (e.g., set fee)
+    await pythTest.sendExecuteGovernanceAction(
+      deployer.getSender(),
+      Buffer.from(PYTH_SET_FEE, "hex")
+    );
 
-    // TODO: add more tests for other governance sequences
+    // Check that the sequence has increased
+    result = await pythTest.getLastExecutedGovernanceSequence();
+    expect(result).toEqual(1);
   });
 
-  it("should return the correct governance data source index", async () => {
-    await deployContract();
+  it("should correctly get governance data source index", async () => {
+    // Deploy contract with initial governance data source
+    await deployContract(
+      BTC_PRICE_FEED_ID,
+      TIME_PERIOD,
+      PRICE,
+      EMA_PRICE,
+      SINGLE_UPDATE_FEE,
+      DATA_SOURCES,
+      0,
+      [TEST_GUARDIAN_ADDRESS1],
+      60051,
+      1,
+      "0000000000000000000000000000000000000000000000000000000000000004",
+      TEST_GOVERNANCE_DATA_SOURCES[0]
+    );
+
+    // Check initial value
+    let result = await pythTest.getGovernanceDataSourceIndex();
+    expect(result).toEqual(0);
 
-    const result = await pythTest.getGovernanceDataSourceIndex();
-    expect(result).toEqual(0); // Initial value should be 0
+    // Execute governance action to change data source
+    await pythTest.sendExecuteGovernanceAction(
+      deployer.getSender(),
+      Buffer.from(PYTH_AUTHORIZE_GOVERNANCE_DATA_SOURCE_TRANSFER, "hex")
+    );
 
-    // TODO: add more tests for other governance data source index
+    // Check that the index has increased
+    result = await pythTest.getGovernanceDataSourceIndex();
+    expect(result).toEqual(1);
   });
 
-  it("should return an empty cell for governance data source", async () => {
+  it("should correctly get governance data source", async () => {
+    // Deploy contract without initial governance data source
     await deployContract();
 
-    const result = await pythTest.getGovernanceDataSource();
-    // assert that the result is an empty cell initally
+    // Check initial value (should be empty)
+    let result = await pythTest.getGovernanceDataSource();
     expect(result).toBeDefined();
     expect(result.bits.length).toBe(0);
     expect(result.refs.length).toBe(0);
 
-    // TODO: add more tests for other governance data source
+    // Deploy contract with initial governance data source
+    await deployContract(
+      BTC_PRICE_FEED_ID,
+      TIME_PERIOD,
+      PRICE,
+      EMA_PRICE,
+      SINGLE_UPDATE_FEE,
+      DATA_SOURCES,
+      0,
+      [TEST_GUARDIAN_ADDRESS1],
+      60051,
+      1,
+      "0000000000000000000000000000000000000000000000000000000000000004",
+      TEST_GOVERNANCE_DATA_SOURCES[0]
+    );
+
+    // Check that the governance data source is set
+    result = await pythTest.getGovernanceDataSource();
+    let dataSource = parseDataSource(result);
+    expect(dataSource).toEqual(TEST_GOVERNANCE_DATA_SOURCES[0]);
+
+    // Execute governance action to change data source
+    await pythTest.sendExecuteGovernanceAction(
+      deployer.getSender(),
+      Buffer.from(PYTH_AUTHORIZE_GOVERNANCE_DATA_SOURCE_TRANSFER, "hex")
+    );
+
+    // Check that the data source has changed
+    result = await pythTest.getGovernanceDataSource();
+    dataSource = parseDataSource(result);
+    expect(dataSource).toEqual(TEST_GOVERNANCE_DATA_SOURCES[1]);
+  });
+
+  it("should correctly get single update fee", async () => {
+    await deployContract();
+
+    // Get the initial fee
+    const result = await pythTest.getSingleUpdateFee();
+
+    expect(result).toBe(SINGLE_UPDATE_FEE);
+  });
+
+  it("should execute set fee governance instruction", async () => {
+    await deployContract(
+      BTC_PRICE_FEED_ID,
+      TIME_PERIOD,
+      PRICE,
+      EMA_PRICE,
+      SINGLE_UPDATE_FEE,
+      DATA_SOURCES,
+      0,
+      [TEST_GUARDIAN_ADDRESS1],
+      60051, // CHAIN_ID of starknet since we are using the test payload for starknet
+      1,
+      "0000000000000000000000000000000000000000000000000000000000000004",
+      TEST_GOVERNANCE_DATA_SOURCES[0]
+    );
+
+    // Execute the governance action
+    const result = await pythTest.sendExecuteGovernanceAction(
+      deployer.getSender(),
+      Buffer.from(PYTH_SET_DATA_SOURCES, "hex")
+    );
+    expect(result.transactions).toHaveTransaction({
+      from: deployer.address,
+      to: pythTest.address,
+      success: true,
+    });
+
+    // Verify that the new data sources are set correctly
+    const newDataSources: DataSource[] = [
+      {
+        emitterChain: 1,
+        emitterAddress:
+          "6bb14509a612f01fbbc4cffeebd4bbfb492a86df717ebe92eb6df432a3f00a25",
+      },
+      {
+        emitterChain: 3,
+        emitterAddress:
+          "000000000000000000000000000000000000000000000000000000000000012d",
+      },
+    ];
+
+    for (const dataSource of newDataSources) {
+      const isValid = await pythTest.getIsValidDataSource(dataSource);
+      expect(isValid).toBe(true);
+    }
+
+    // Verify that the old data source is no longer valid
+    const oldDataSource = DATA_SOURCES[0];
+    const oldDataSourceIsValid = await pythTest.getIsValidDataSource(
+      oldDataSource
+    );
+    expect(oldDataSourceIsValid).toBe(false);
+  });
+
+  it("should execute authorize governance data source transfer", async () => {
+    await deployContract(
+      BTC_PRICE_FEED_ID,
+      TIME_PERIOD,
+      PRICE,
+      EMA_PRICE,
+      SINGLE_UPDATE_FEE,
+      DATA_SOURCES,
+      0,
+      [TEST_GUARDIAN_ADDRESS1],
+      60051, // CHAIN_ID of starknet since we are using the test payload for starknet
+      1,
+      "0000000000000000000000000000000000000000000000000000000000000004",
+      TEST_GOVERNANCE_DATA_SOURCES[0]
+    );
+
+    // Get the initial governance data source index
+    const initialIndex = await pythTest.getGovernanceDataSourceIndex();
+    expect(initialIndex).toEqual(0); // Initial value should be 0
+
+    // Get the initial governance data source
+    const initialDataSourceCell = await pythTest.getGovernanceDataSource();
+    const initialDataSource = parseDataSource(initialDataSourceCell);
+    expect(initialDataSource).toEqual(TEST_GOVERNANCE_DATA_SOURCES[0]);
+
+    // Get the initial last executed governance sequence
+    const initialSequence = await pythTest.getLastExecutedGovernanceSequence();
+    expect(initialSequence).toEqual(0); // Initial value should be 0
+
+    // Execute the governance action
+    const result = await pythTest.sendExecuteGovernanceAction(
+      deployer.getSender(),
+      Buffer.from(PYTH_AUTHORIZE_GOVERNANCE_DATA_SOURCE_TRANSFER, "hex")
+    );
+    expect(result.transactions).toHaveTransaction({
+      from: deployer.address,
+      to: pythTest.address,
+      success: true,
+    });
+
+    // Get the new governance data source index
+    const newIndex = await pythTest.getGovernanceDataSourceIndex();
+    expect(newIndex).toEqual(1); // The new index value should match the one in the test payload
+
+    // Get the new governance data source
+    const newDataSourceCell = await pythTest.getGovernanceDataSource();
+    const newDataSource = parseDataSource(newDataSourceCell);
+    expect(newDataSource).not.toEqual(initialDataSource); // The data source should have changed
+    expect(newDataSource).toEqual(TEST_GOVERNANCE_DATA_SOURCES[1]); // The data source should have changed
+
+    // Get the new last executed governance sequence
+    const newSequence = await pythTest.getLastExecutedGovernanceSequence();
+    expect(newSequence).toBeGreaterThan(initialSequence); // The sequence should have increased
+    expect(newSequence).toBe(1);
+  });
+
+  it("should fail when executing request governance data source transfer directly", async () => {
+    await deployContract(
+      BTC_PRICE_FEED_ID,
+      TIME_PERIOD,
+      PRICE,
+      EMA_PRICE,
+      SINGLE_UPDATE_FEE,
+      DATA_SOURCES,
+      0,
+      [TEST_GUARDIAN_ADDRESS1],
+      60051, // CHAIN_ID of starknet since we are using the test payload for starknet
+      1,
+      "0000000000000000000000000000000000000000000000000000000000000004",
+      TEST_GOVERNANCE_DATA_SOURCES[1]
+    );
+
+    const result = await pythTest.sendExecuteGovernanceAction(
+      deployer.getSender(),
+      Buffer.from(PYTH_REQUEST_GOVERNANCE_DATA_SOURCE_TRANSFER, "hex")
+    );
+
+    // Check that the transaction did not succeed
+    expect(result.transactions).toHaveTransaction({
+      from: deployer.address,
+      to: pythTest.address,
+      success: false,
+      exitCode: 1012, // ERROR_INVALID_GOVERNANCE_ACTION = 1012
+    });
+
+    // Verify that the governance data source index hasn't changed
+    const index = await pythTest.getGovernanceDataSourceIndex();
+    expect(index).toEqual(0); // Should still be the initial value
+
+    // Verify that the governance data source hasn't changed
+    const dataSourceCell = await pythTest.getGovernanceDataSource();
+    const dataSource = parseDataSource(dataSourceCell);
+    expect(dataSource).toEqual(TEST_GOVERNANCE_DATA_SOURCES[1]); // Should still be the initial value
+  });
+
+  it("should fail to execute governance action with invalid governance data source", async () => {
+    await deployContract(
+      BTC_PRICE_FEED_ID,
+      TIME_PERIOD,
+      PRICE,
+      EMA_PRICE,
+      SINGLE_UPDATE_FEE,
+      DATA_SOURCES,
+      0,
+      [TEST_GUARDIAN_ADDRESS1],
+      60051,
+      1,
+      "0000000000000000000000000000000000000000000000000000000000000004",
+      TEST_GOVERNANCE_DATA_SOURCES[1]
+    );
+
+    const result = await pythTest.sendExecuteGovernanceAction(
+      deployer.getSender(),
+      Buffer.from(PYTH_SET_FEE, "hex")
+    );
+
+    expect(result.transactions).toHaveTransaction({
+      from: deployer.address,
+      to: pythTest.address,
+      success: false,
+      exitCode: 1032, // ERROR_INVALID_GOVERNANCE_DATA_SOURCE
+    });
+  });
+
+  it("should fail to execute governance action with old sequence number", async () => {
+    await deployContract(
+      BTC_PRICE_FEED_ID,
+      TIME_PERIOD,
+      PRICE,
+      EMA_PRICE,
+      SINGLE_UPDATE_FEE,
+      DATA_SOURCES,
+      0,
+      [TEST_GUARDIAN_ADDRESS1],
+      60051,
+      1,
+      "0000000000000000000000000000000000000000000000000000000000000004",
+      TEST_GOVERNANCE_DATA_SOURCES[0]
+    );
+
+    // Execute a governance action to increase the sequence number
+    await pythTest.sendExecuteGovernanceAction(
+      deployer.getSender(),
+      Buffer.from(PYTH_SET_FEE, "hex")
+    );
+
+    // Try to execute the same governance action again
+    const result = await pythTest.sendExecuteGovernanceAction(
+      deployer.getSender(),
+      Buffer.from(PYTH_SET_FEE, "hex")
+    );
+
+    expect(result.transactions).toHaveTransaction({
+      from: deployer.address,
+      to: pythTest.address,
+      success: false,
+      exitCode: 1033, // ERROR_OLD_GOVERNANCE_MESSAGE
+    });
   });
 });

target_chains/ton/contracts/tests/utils.ts

@@ -1,4 +1,5 @@
-import { Cell, beginCell } from "@ton/core";
+import { DataSource } from "@pythnetwork/xc-admin-common";
+import { Cell, Transaction, beginCell } from "@ton/core";
 
 export function createCellChain(buffer: Buffer): Cell {
   let chunks = bufferToChunks(buffer, 127);
@@ -43,3 +44,26 @@ function bufferToChunks(
 
   return chunks;
 }
+
+// Helper function to parse DataSource from a Cell
+export function parseDataSource(cell: Cell): DataSource {
+  const slice = cell.beginParse();
+  const emitterChain = slice.loadUint(16);
+  const emitterAddress = slice.loadUint(256).toString(16).padStart(64, "0");
+  return { emitterChain, emitterAddress };
+}
+
+function computedGeneric(transaction: Transaction) {
+  if (transaction.description.type !== "generic")
+    throw "Expected generic transactionaction";
+  if (transaction.description.computePhase.type !== "vm")
+    throw "Compute phase expected";
+  return transaction.description.computePhase;
+}
+
+export function printTxGasStats(name: string, transaction: Transaction) {
+  const txComputed = computedGeneric(transaction);
+  console.log(`${name} used ${txComputed.gasUsed} gas`);
+  console.log(`${name} gas cost: ${txComputed.gasFees}`);
+  return txComputed.gasFees;
+}

target_chains/ton/contracts/tests/utils/pyth.ts

@@ -36,3 +36,26 @@ export const HERMES_BTC_ETH_UPDATE =
 
 export const BTC_PRICE_FEED_ID =
   "0xe62df6c8b4a85fe1a67db44dc12de5db330f7ac66b72dc658afedf0f4a415b43";
+
+export const TEST_GUARDIAN_ADDRESS1 =
+  "0x686b9ea8e3237110eaaba1f1b7467559a3273819";
+
+// A Pyth governance instruction to authorize governance data source transfer signed by the test guardian #1.
+// From: target_chains/starknet/contracts/tests/data.cairo::pyth_auth_transfer()
+export const PYTH_AUTHORIZE_GOVERNANCE_DATA_SOURCE_TRANSFER =
+  "01000000000100e4c6595b44ed764ebf9d563e8b2e8233cc24f7c35737e83c4ca1ec51f77dfd6214a146fa57420f97d51e7161342b4833b8e75c89a3895e609d7d58da7ffb5b1a000000000100000002000100000000000000000000000000000000000000000000000000000000000000290000000000000001065054474d0101ea9301000000000100b3abee2eed7d504284c57387abcbb87fe9cf4807228d2d08b776ea94347bdaa73e6958ce95a49d11b9a07fdc93a6705b666fb8a76c0cff0578eb2c881d80b29e0100000001000000020002000000000000000000000000000000000000000000000000000000000000002b0000000000000001065054474d0105ea9300000001";
+
+// A Pyth governance instruction to set fee signed by the test guardian #1.
+// From: target_chains/starknet/contracts/tests/data.cairo::pyth_set_fee()
+export const PYTH_SET_FEE =
+  "010000000001006da27b990a357166853242ffec67013c89696f82d009ce79b6cb302db14f2e2e3ec3513c47ce572524ac42fedd7fb4100303baafd9ad5de6e7ed587713a36a2b010000000100000002000100000000000000000000000000000000000000000000000000000000000000290000000000000001065054474d0103ea93000000000000002a0000000000000002";
+
+// A Pyth governance instruction to set data sources signed by the test guardian #1.
+// From: target_chains/starknet/contracts/tests/data.cairo::pyth_set_data_sources()
+export const PYTH_SET_DATA_SOURCES =
+  "01000000000100671d487654ad77101243588c74a9f9d90de187b9807445f9b4b0bc2eb3363b1d72aff4ad4f80a09e6cdd84e29b2e513a50efc66c979beef21ca5095d425fa9df000000000100000002000100000000000000000000000000000000000000000000000000000000000000290000000000000001065054474d0102ea930200016bb14509a612f01fbbc4cffeebd4bbfb492a86df717ebe92eb6df432a3f00a250003000000000000000000000000000000000000000000000000000000000000012d";
+
+// A Pyth governance instruction to request governance data source transfer signed by the test guardian #1.
+// From: target_chains/starknet/contracts/tests/data.cairo::pyth_request_transfer()
+export const PYTH_REQUEST_GOVERNANCE_DATA_SOURCE_TRANSFER =
+  "01000000000100b3abee2eed7d504284c57387abcbb87fe9cf4807228d2d08b776ea94347bdaa73e6958ce95a49d11b9a07fdc93a6705b666fb8a76c0cff0578eb2c881d80b29e0100000001000000020002000000000000000000000000000000000000000000000000000000000000002b0000000000000001065054474d0105ea9300000001";

target_chains/ton/contracts/wrappers/PythTest.ts

@@ -27,6 +27,7 @@ export type PythTestConfig = {
   chainId: number;
   governanceChainId: number;
   governanceContract: string;
+  governanceDataSource?: DataSource;
 };
 
 export class PythTest implements Contract {
@@ -51,7 +52,8 @@ export class PythTest implements Contract {
       config.guardianSet,
       config.chainId,
       config.governanceChainId,
-      config.governanceContract
+      config.governanceContract,
+      config.governanceDataSource
     );
     const init = { code, data };
     return new PythTest(contractAddress(workchain, init), init);
@@ -68,7 +70,8 @@ export class PythTest implements Contract {
     guardianSet: string[],
     chainId: number,
     governanceChainId: number,
-    governanceContract: string
+    governanceContract: string,
+    governanceDataSource?: DataSource
   ): Cell {
     const priceDict = Dictionary.empty(
       Dictionary.Keys.BigUint(256),
@@ -143,7 +146,16 @@ export class PythTest implements Contract {
       .storeUint(governanceChainId, 16)
       .storeBuffer(Buffer.from(governanceContract, "hex"))
       .storeDict(Dictionary.empty()) // consumed_governance_actions
-      .storeRef(beginCell()) // governance_data_source, empty for initial state
+      .storeRef(
+        governanceDataSource
+          ? beginCell()
+              .storeUint(governanceDataSource.emitterChain, 16)
+              .storeBuffer(
+                Buffer.from(governanceDataSource.emitterAddress, "hex")
+              )
+              .endCell()
+          : beginCell().endCell()
+      ) // governance_data_source
       .storeUint(0, 64) // last_executed_governance_sequence
       .storeUint(0, 32) // governance_data_source_index
       .endCell();
@@ -255,6 +267,11 @@ export class PythTest implements Contract {
     return result.stack.readNumber();
   }
 
+  async getSingleUpdateFee(provider: ContractProvider) {
+    const result = await provider.get("test_get_single_update_fee", []);
+    return result.stack.readNumber();
+  }
+
   async sendUpdatePriceFeeds(
     provider: ContractProvider,
     via: Sender,
@@ -315,4 +332,37 @@ export class PythTest implements Contract {
     const result = await provider.get("test_get_governance_data_source", []);
     return result.stack.readCell();
   }
+
+  async sendExecuteGovernanceAction(
+    provider: ContractProvider,
+    via: Sender,
+    governanceAction: Buffer
+  ) {
+    const messageBody = beginCell()
+      .storeUint(3, 32) // OP_EXECUTE_GOVERNANCE_ACTION
+      .storeRef(createCellChain(governanceAction))
+      .endCell();
+
+    await provider.internal(via, {
+      value: toNano("0.1"),
+      sendMode: SendMode.PAY_GAS_SEPARATELY,
+      body: messageBody,
+    });
+  }
+
+  async getIsValidDataSource(
+    provider: ContractProvider,
+    dataSource: DataSource
+  ) {
+    const result = await provider.get("test_get_is_valid_data_source", [
+      {
+        type: "cell",
+        cell: beginCell()
+          .storeUint(dataSource.emitterChain, 16)
+          .storeUint(BigInt("0x" + dataSource.emitterAddress), 256)
+          .endCell(),
+      },
+    ]);
+    return result.stack.readBoolean();
+  }
 }