TEST: Restaking Variations (#41)

Added tests for:
+ if operator is "removed" from the NCN
+ if vault is "removed" from the NCN
+ if vault is not updated

clients/js/jito_tip_router/errors/jitoTipRouter.ts

@@ -160,10 +160,12 @@ export const JITO_TIP_ROUTER_ERROR__ROUTER_STILL_ROUTING = 0x223d; // 8765
 export const JITO_TIP_ROUTER_ERROR__INVALID_EPOCHS_BEFORE_STALL = 0x223e; // 8766
 /** InvalidSlotsAfterConsensus: Invalid slots after consensus */
 export const JITO_TIP_ROUTER_ERROR__INVALID_SLOTS_AFTER_CONSENSUS = 0x223f; // 8767
+/** VaultNeedsUpdate: Vault needs to be updated */
+export const JITO_TIP_ROUTER_ERROR__VAULT_NEEDS_UPDATE = 0x2240; // 8768
 /** InvalidAccountStatus: Invalid Account Status */
-export const JITO_TIP_ROUTER_ERROR__INVALID_ACCOUNT_STATUS = 0x2240; // 8768
+export const JITO_TIP_ROUTER_ERROR__INVALID_ACCOUNT_STATUS = 0x2241; // 8769
 /** AccountAlreadyInitialized: Account already initialized */
-export const JITO_TIP_ROUTER_ERROR__ACCOUNT_ALREADY_INITIALIZED = 0x2241; // 8769
+export const JITO_TIP_ROUTER_ERROR__ACCOUNT_ALREADY_INITIALIZED = 0x2242; // 8770
 
 export type JitoTipRouterError =
   | typeof JITO_TIP_ROUTER_ERROR__ACCOUNT_ALREADY_INITIALIZED
@@ -229,6 +231,7 @@ export type JitoTipRouterError =
   | typeof JITO_TIP_ROUTER_ERROR__TOO_MANY_VAULT_OPERATOR_DELEGATIONS
   | typeof JITO_TIP_ROUTER_ERROR__TOO_MANY_VAULTS_FOR_REGISTRY
   | typeof JITO_TIP_ROUTER_ERROR__VAULT_INDEX_ALREADY_IN_USE
+  | typeof JITO_TIP_ROUTER_ERROR__VAULT_NEEDS_UPDATE
   | typeof JITO_TIP_ROUTER_ERROR__VAULT_NOT_IN_REGISTRY
   | typeof JITO_TIP_ROUTER_ERROR__VAULT_OPERATOR_DELEGATION_FINALIZED
   | typeof JITO_TIP_ROUTER_ERROR__VAULT_REGISTRY_LIST_FULL
@@ -308,6 +311,7 @@ if (process.env.NODE_ENV !== 'production') {
     [JITO_TIP_ROUTER_ERROR__TOO_MANY_VAULT_OPERATOR_DELEGATIONS]: `Too many vault operator delegations`,
     [JITO_TIP_ROUTER_ERROR__TOO_MANY_VAULTS_FOR_REGISTRY]: `Too many vaults for registry`,
     [JITO_TIP_ROUTER_ERROR__VAULT_INDEX_ALREADY_IN_USE]: `Vault index already in use by a different mint`,
+    [JITO_TIP_ROUTER_ERROR__VAULT_NEEDS_UPDATE]: `Vault needs to be updated`,
     [JITO_TIP_ROUTER_ERROR__VAULT_NOT_IN_REGISTRY]: `Vault not in weight table registry`,
     [JITO_TIP_ROUTER_ERROR__VAULT_OPERATOR_DELEGATION_FINALIZED]: `Vault operator delegation is already finalized - should not happen`,
     [JITO_TIP_ROUTER_ERROR__VAULT_REGISTRY_LIST_FULL]: `Vault Registry mints are at capacity`,

clients/rust/jito_tip_router/src/generated/errors/jito_tip_router.rs

@@ -229,12 +229,15 @@ pub enum JitoTipRouterError {
     /// 8767 - Invalid slots after consensus
     #[error("Invalid slots after consensus")]
     InvalidSlotsAfterConsensus = 0x223F,
-    /// 8768 - Invalid Account Status
+    /// 8768 - Vault needs to be updated
+    #[error("Vault needs to be updated")]
+    VaultNeedsUpdate = 0x2240,
+    /// 8769 - Invalid Account Status
     #[error("Invalid Account Status")]
-    InvalidAccountStatus = 0x2240,
-    /// 8769 - Account already initialized
+    InvalidAccountStatus = 0x2241,
+    /// 8770 - Account already initialized
     #[error("Account already initialized")]
-    AccountAlreadyInitialized = 0x2241,
+    AccountAlreadyInitialized = 0x2242,
 }
 
 impl solana_program::program_error::PrintProgramError for JitoTipRouterError {

core/src/error.rs

@@ -151,6 +151,8 @@ pub enum TipRouterError {
     InvalidEpochsBeforeStall,
     #[error("Invalid slots after consensus")]
     InvalidSlotsAfterConsensus,
+    #[error("Vault needs to be updated")]
+    VaultNeedsUpdate,
     #[error("Invalid Account Status")]
     InvalidAccountStatus,
     #[error("Account already initialized")]

idl/jito_tip_router.json

@@ -3720,11 +3720,16 @@
     },
     {
       "code": 8768,
+      "name": "VaultNeedsUpdate",
+      "msg": "Vault needs to be updated"
+    },
+    {
+      "code": 8769,
       "name": "InvalidAccountStatus",
       "msg": "Invalid Account Status"
     },
     {
-      "code": 8769,
+      "code": 8770,
       "name": "AccountAlreadyInitialized",
       "msg": "Account already initialized"
     }

integration_tests/tests/fixtures/test_builder.rs

@@ -19,6 +19,7 @@ use solana_program::{
 use solana_program_test::{processor, BanksClientError, ProgramTest, ProgramTestContext};
 use solana_sdk::{
     account::Account,
+    clock::DEFAULT_SLOTS_PER_EPOCH,
     commitment_config::CommitmentLevel,
     epoch_schedule::EpochSchedule,
     native_token::lamports_to_sol,
@@ -188,6 +189,22 @@ impl TestBuilder {
         Ok(())
     }
 
+    pub async fn warp_epoch_incremental(
+        &mut self,
+        incremental_epochs: u64,
+    ) -> Result<(), BanksClientError> {
+        let clock: Clock = self.context.banks_client.get_sysvar().await?;
+        self.context
+            .warp_to_slot(
+                clock
+                    .slot
+                    .checked_add(DEFAULT_SLOTS_PER_EPOCH * incremental_epochs)
+                    .unwrap(),
+            )
+            .map_err(|_| BanksClientError::ClientError("failed to warp slot"))?;
+        Ok(())
+    }
+
     pub async fn set_account(&mut self, address: Pubkey, account: Account) {
         self.context
             .borrow_mut()
@@ -692,16 +709,44 @@ impl TestBuilder {
         test_ncn: &TestNcn,
     ) -> TestResult<()> {
         let mut tip_router_client = self.tip_router_client();
+        let mut vault_program_client = self.vault_program_client();
 
         let clock = self.clock().await;
+        let slot = clock.slot;
         let epoch = clock.epoch;
         let ncn = test_ncn.ncn_root.ncn_pubkey;
 
+        let operators_for_update = test_ncn
+            .operators
+            .iter()
+            .map(|operator_root| operator_root.operator_pubkey)
+            .collect::<Vec<Pubkey>>();
+
         for operator_root in test_ncn.operators.iter() {
             let operator = operator_root.operator_pubkey;
+
+            let operator_snapshot = tip_router_client
+                .get_operator_snapshot(operator, ncn, epoch)
+                .await?;
+
+            // If operator snapshot is finalized it means that the operator is not active.
+            if operator_snapshot.finalized() {
+                continue;
+            }
+
             for vault_root in test_ncn.vaults.iter() {
                 let vault = vault_root.vault_pubkey;
 
+                let vault_is_update_needed = vault_program_client
+                    .get_vault_is_update_needed(&vault, slot)
+                    .await?;
+
+                if vault_is_update_needed {
+                    vault_program_client
+                        .do_full_vault_update(&vault, &operators_for_update)
+                        .await?;
+                }
+
                 tip_router_client
                     .do_snapshot_vault_operator_delegation(vault, operator, ncn, epoch)
                     .await?;
@@ -852,19 +897,6 @@ impl TestBuilder {
                         continue;
                     }
 
-                    //TODO remove
-                    // let (ncn_reward_receiver, _, _) = NcnRewardReceiver::find_program_address(
-                    //     &jito_tip_router_program::id(),
-                    //     *group,
-                    //     &operator,
-                    //     &ncn,
-                    //     epoch,
-                    // );
-                    // let sol_rewards = lamports_to_sol(rewards);
-                    // tip_router_client
-                    //     .airdrop(&ncn_reward_receiver, sol_rewards)
-                    //     .await?;
-
                     tip_router_client
                         .do_distribute_base_ncn_reward_route(*group, operator, ncn, epoch)
                         .await?;

integration_tests/tests/fixtures/vault_client.rs

@@ -222,6 +222,20 @@ impl VaultProgramClient {
         )?)
     }
 
+    pub async fn get_vault_is_update_needed(
+        &mut self,
+        vault: &Pubkey,
+        slot: u64,
+    ) -> Result<bool, TestError> {
+        let vault_config = self
+            .get_config(&Config::find_program_address(&jito_vault_program::id()).0)
+            .await?;
+        let vault_account = self.get_vault(vault).await?;
+
+        let is_update_needed = vault_account.is_update_needed(slot, vault_config.epoch_length())?;
+        Ok(is_update_needed)
+    }
+
     pub async fn do_initialize_config(&mut self) -> Result<Keypair, TestError> {
         let config_admin = Keypair::new();
 

integration_tests/tests/tip_router/epoch_state.rs

@@ -10,14 +10,11 @@ mod tests {
     #[tokio::test]
     async fn test_all_test_ncn_functions_pt1() -> TestResult<()> {
         let mut fixture = TestBuilder::new().await;
-        let mut stake_pool_client = fixture.stake_pool_client();
         let mut tip_router_client = fixture.tip_router_client();
 
         const OPERATOR_COUNT: usize = 2;
         const VAULT_COUNT: usize = 3;
 
-        let pool_root = stake_pool_client.do_initialize_stake_pool().await?;
-
         let test_ncn = fixture
             .create_initial_test_ncn(OPERATOR_COUNT, VAULT_COUNT, Some(100))
             .await?;
@@ -67,6 +64,46 @@ mod tests {
             }
         }
 
+        // To be continued... Running into stack overflow issues
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_all_test_ncn_functions_pt2() -> TestResult<()> {
+        let mut fixture = TestBuilder::new().await;
+        let mut stake_pool_client = fixture.stake_pool_client();
+        let mut tip_router_client = fixture.tip_router_client();
+
+        const OPERATOR_COUNT: usize = 2;
+        const VAULT_COUNT: usize = 3;
+
+        let pool_root = stake_pool_client.do_initialize_stake_pool().await?;
+
+        let test_ncn = fixture
+            .create_initial_test_ncn(OPERATOR_COUNT, VAULT_COUNT, Some(100))
+            .await?;
+        let ncn = test_ncn.ncn_root.ncn_pubkey;
+        let epoch = fixture.clock().await.epoch;
+
+        {
+            fixture.add_epoch_state_for_test_ncn(&test_ncn).await?;
+        }
+
+        {
+            fixture.add_admin_weights_for_test_ncn(&test_ncn).await?;
+        }
+
+        {
+            fixture.add_epoch_snapshot_to_test_ncn(&test_ncn).await?;
+        }
+
+        {
+            fixture
+                .add_operator_snapshots_to_test_ncn(&test_ncn)
+                .await?;
+        }
+
         {
             fixture
                 .add_vault_operator_delegation_snapshots_to_test_ncn(&test_ncn)
@@ -123,22 +160,13 @@ mod tests {
             }
         }
 
-        {
-            fixture
-                .route_in_base_rewards_for_test_ncn(&test_ncn, 10_000, &pool_root)
-                .await?;
-            fixture
-                .route_in_ncn_rewards_for_test_ncn(&test_ncn, &pool_root)
-                .await?;
-        }
-
         // To be continued... Running into stack overflow issues
 
         Ok(())
     }
 
     #[tokio::test]
-    async fn test_all_test_ncn_functions_pt2() -> TestResult<()> {
+    async fn test_all_test_ncn_functions_pt3() -> TestResult<()> {
         let mut fixture = TestBuilder::new().await;
         let mut stake_pool_client = fixture.stake_pool_client();
         let mut tip_router_client = fixture.tip_router_client();

integration_tests/tests/tip_router/mod.rs

@@ -15,6 +15,7 @@ mod initialize_vault_registry;
 mod initialize_weight_table;
 mod meta_tests;
 mod register_vault;
+mod restaking_variations;
 mod set_config_fees;
 mod set_new_admin;
 mod set_tie_breaker;