test: test aggregation toggle

Author: Reisen

program/rust/src/tests/test_aggregate_v2.rs

@@ -6,96 +6,173 @@ use {
             PriceAccountFlags,
             PythAccount,
         },
-        c_oracle_header::PC_VERSION,
-        deserialize::load_checked,
+        c_oracle_header::{
+            PC_STATUS_TRADING,
+            PC_VERSION,
+        },
+        deserialize::{
+            load_checked,
+            load_mut,
+        },
         instruction::{
             AddPublisherArgs,
             OracleCommand,
+            UpdPriceArgs,
         },
         processor::{
             process_instruction,
+            DISABLE_ACCUMULATOR_V2,
             ENABLE_ACCUMULATOR_V2,
         },
-        tests::test_utils::AccountSetup,
+        tests::test_utils::{
+            update_clock_slot,
+            AccountSetup,
+        },
     },
     bytemuck::bytes_of,
     solana_program::{
         pubkey::Pubkey,
         rent::Rent,
     },
+    solana_sdk::account_info::AccountInfo,
+    std::mem::size_of,
 };
 
-#[test]
-fn test_aggregate_v2() {
-    let program_id = Pubkey::new_unique();
-    let publisher = Pubkey::new_unique();
-
-    let mut cmd = AddPublisherArgs {
-        header: OracleCommand::AddPublisher.into(),
-        publisher,
-    };
-    let mut instruction_data = bytes_of::<AddPublisherArgs>(&cmd);
-
-    let mut funding_setup = AccountSetup::new_funding();
-    let funding_account = funding_setup.as_account_info();
+struct Accounts {
+    program_id:          Pubkey,
+    publisher_account:   AccountSetup,
+    funding_account:     AccountSetup,
+    price_account:       AccountSetup,
+    permissions_account: AccountSetup,
+    clock_account:       AccountSetup,
+}
 
-    let mut price_setup = AccountSetup::new::<PriceAccount>(&program_id);
-    let price_account = price_setup.as_account_info();
-    PriceAccount::initialize(&price_account, PC_VERSION).unwrap();
+impl Accounts {
+    fn new() -> Self {
+        let program_id = Pubkey::new_unique();
+        let publisher_account = AccountSetup::new_funding();
+        let clock_account = AccountSetup::new_clock();
+        let mut funding_account = AccountSetup::new_funding();
+        let mut permissions_account = AccountSetup::new_permission(&program_id);
+        let mut price_account = AccountSetup::new::<PriceAccount>(&program_id);
 
-    **price_account.try_borrow_mut_lamports().unwrap() = 100;
+        PriceAccount::initialize(&price_account.as_account_info(), PC_VERSION).unwrap();
 
-    let mut permissions_setup = AccountSetup::new_permission(&program_id);
-    let permissions_account = permissions_setup.as_account_info();
+        {
+            let permissions_account_info = permissions_account.as_account_info();
+            let mut permissions_account_data =
+                PermissionAccount::initialize(&permissions_account_info, PC_VERSION).unwrap();
+            permissions_account_data.master_authority = *funding_account.as_account_info().key;
+            permissions_account_data.data_curation_authority =
+                *funding_account.as_account_info().key;
+            permissions_account_data.security_authority = *funding_account.as_account_info().key;
+        }
 
-    {
-        let mut permissions_account_data =
-            PermissionAccount::initialize(&permissions_account, PC_VERSION).unwrap();
-        permissions_account_data.master_authority = *funding_account.key;
-        permissions_account_data.data_curation_authority = *funding_account.key;
-        permissions_account_data.security_authority = *funding_account.key;
+        Self {
+            program_id,
+            publisher_account,
+            funding_account,
+            price_account,
+            permissions_account,
+            clock_account,
+        }
     }
+}
 
-    // Give the price account enough lamports to be rent exempt
-    **price_account.try_borrow_mut_lamports().unwrap() =
-        Rent::minimum_balance(&Rent::default(), PriceAccount::MINIMUM_SIZE);
+fn add_publisher(accounts: &mut Accounts, publisher: Option<Pubkey>) {
+    let args = AddPublisherArgs {
+        header:    OracleCommand::AddPublisher.into(),
+        publisher: publisher.unwrap_or(*accounts.publisher_account.as_account_info().key),
+    };
 
     assert!(process_instruction(
-        &program_id,
+        &accounts.program_id,
         &[
-            funding_account.clone(),
-            price_account.clone(),
-            permissions_account.clone(),
+            accounts.funding_account.as_account_info(),
+            accounts.price_account.as_account_info(),
+            accounts.permissions_account.as_account_info(),
         ],
-        instruction_data
+        bytes_of::<AddPublisherArgs>(&args)
     )
     .is_ok());
+}
 
-    {
-        let price_data = load_checked::<PriceAccount>(&price_account, PC_VERSION).unwrap();
-        assert_eq!(price_data.num_, 1);
-        assert_eq!(price_data.header.size, PriceAccount::INITIAL_SIZE);
-        assert!(price_data.comp_[0].pub_ == publisher);
-        // Make sure that v2 aggregation is disabled
-        assert!(!price_data.flags.contains(PriceAccountFlags::ACCUMULATOR_V2));
-    }
+fn update_price(accounts: &mut Accounts, price: i64, conf: u64, slot: u64) {
+    let instruction_data = &mut [0u8; size_of::<UpdPriceArgs>()];
+    let mut cmd = load_mut::<UpdPriceArgs>(instruction_data).unwrap();
+    cmd.header = OracleCommand::UpdPrice.into();
+    cmd.status = PC_STATUS_TRADING;
+    cmd.price = price;
+    cmd.confidence = conf;
+    cmd.publishing_slot = slot;
+    cmd.unused_ = 0;
 
-    cmd.publisher = ENABLE_ACCUMULATOR_V2.into();
-    instruction_data = bytes_of::<AddPublisherArgs>(&cmd);
-    assert!(process_instruction(
-        &program_id,
+    let mut clock = accounts.clock_account.as_account_info();
+    clock.is_signer = false;
+    clock.is_writable = false;
+
+    process_instruction(
+        &accounts.program_id,
         &[
-            funding_account.clone(),
-            price_account.clone(),
-            permissions_account.clone(),
+            accounts.publisher_account.as_account_info(),
+            accounts.price_account.as_account_info(),
+            clock,
         ],
-        instruction_data
+        instruction_data,
     )
-    .is_ok());
+    .unwrap();
+}
+
+#[test]
+fn test_aggregate_v2_toggle() {
+    let accounts = &mut Accounts::new();
+
+    // Add an initial Publisher to test with.
+    add_publisher(accounts, None);
+
+    // Update the price, no aggregation will happen on the first slot.
+    {
+        update_clock_slot(&mut accounts.clock_account.as_account_info(), 1);
+        update_price(accounts, 42, 2, 1);
+        let info = accounts.price_account.as_account_info();
+        let price_data = load_checked::<PriceAccount>(&info, PC_VERSION).unwrap();
+        assert_eq!(price_data.last_slot_, 0);
+        assert!(!price_data.flags.contains(PriceAccountFlags::ACCUMULATOR_V2));
+    }
+
+    // Update again, component is now TRADING so aggregation should trigger.
+    {
+        update_clock_slot(&mut accounts.clock_account.as_account_info(), 2);
+        update_price(accounts, 42, 2, 2);
+        let info = accounts.price_account.as_account_info();
+        let price_data = load_checked::<PriceAccount>(&info, PC_VERSION).unwrap();
+        assert_eq!(price_data.last_slot_, 2);
+        assert!(!price_data.flags.contains(PriceAccountFlags::ACCUMULATOR_V2));
+    }
 
-    // Make sure that v2 aggregation is enabled
+    // Enable v2 Aggregation
+    add_publisher(accounts, Some(ENABLE_ACCUMULATOR_V2.into()));
+
+    // Update again, with accumulator bit set, aggregation should not have
+    // happened, as its now the validators job.
     {
-        let price_data = load_checked::<PriceAccount>(&price_account, PC_VERSION).unwrap();
+        update_clock_slot(&mut accounts.clock_account.as_account_info(), 3);
+        update_price(accounts, 42, 2, 3);
+        let info = accounts.price_account.as_account_info();
+        let price_data = load_checked::<PriceAccount>(&info, PC_VERSION).unwrap();
+        assert_eq!(price_data.last_slot_, 2);
         assert!(price_data.flags.contains(PriceAccountFlags::ACCUMULATOR_V2));
     }
+
+    add_publisher(accounts, Some(DISABLE_ACCUMULATOR_V2.into()));
+
+    // Confirm disabling v2 Aggregation re-enables the aggregation flow.
+    {
+        update_clock_slot(&mut accounts.clock_account.as_account_info(), 4);
+        update_price(accounts, 42, 2, 4);
+        let info = accounts.price_account.as_account_info();
+        let price_data = load_checked::<PriceAccount>(&info, PC_VERSION).unwrap();
+        assert_eq!(price_data.last_slot_, 4);
+        assert!(!price_data.flags.contains(PriceAccountFlags::ACCUMULATOR_V2));
+    }
 }