working

core/src/base_fee_group.rs

@@ -16,7 +16,7 @@ pub enum BaseFeeGroupType {
     Reserved7 = 0x7,
 }
 
-#[derive(Debug, Clone, Copy, Zeroable, ShankType, Pod)]
+#[derive(Debug, Clone, Copy, Zeroable, ShankType, Pod, PartialEq, Eq)]
 #[repr(C)]
 pub struct BaseFeeGroup {
     pub group: u8,
@@ -218,24 +218,26 @@ mod tests {
             Err(TipRouterError::InvalidNcnFeeGroup)
         ));
     }
-
     #[test]
     fn test_all_groups() {
         let all_groups = BaseFeeGroup::all_groups();
 
         // Verify count matches FEE_GROUP_COUNT
         assert_eq!(all_groups.len(), BaseFeeGroup::FEE_GROUP_COUNT);
 
-        // Verify groups are in correct order
-        for (i, group) in all_groups.iter().enumerate() {
-            assert_eq!(group.group as usize, i);
-        }
+        // Verify groups are in correct order and have expected values
+        let expected_groups = vec![
+            BaseFeeGroup::new(BaseFeeGroupType::DAO),
+            BaseFeeGroup::new(BaseFeeGroupType::Reserved1),
+            BaseFeeGroup::new(BaseFeeGroupType::Reserved2),
+            BaseFeeGroup::new(BaseFeeGroupType::Reserved3),
+            BaseFeeGroup::new(BaseFeeGroupType::Reserved4),
+            BaseFeeGroup::new(BaseFeeGroupType::Reserved5),
+            BaseFeeGroup::new(BaseFeeGroupType::Reserved6),
+            BaseFeeGroup::new(BaseFeeGroupType::Reserved7),
+        ];
 
-        // Verify no duplicates
-        let mut unique_groups: Vec<u8> = all_groups.iter().map(|g| g.group).collect();
-        unique_groups.sort();
-        unique_groups.dedup();
-        assert_eq!(unique_groups.len(), BaseFeeGroup::FEE_GROUP_COUNT);
+        assert_eq!(all_groups, expected_groups);
     }
 
     #[test]

core/src/base_reward_router.rs

@@ -610,3 +610,52 @@ impl BaseRewardRouterRewards {
         self.rewards.into()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use solana_program::pubkey::Pubkey;
+
+    use super::*;
+
+    #[test]
+    fn test_route_incoming_rewards() {
+        let mut router = BaseRewardRouter::new(
+            Pubkey::new_unique(), // ncn
+            1,                    // ncn_epoch
+            1,                    // bump
+            100,                  // slot_created
+        );
+
+        // Initial state checks
+        assert_eq!(router.total_rewards(), 0);
+        assert_eq!(router.reward_pool(), 0);
+        assert_eq!(router.rewards_processed(), 0);
+
+        // Test routing 1000 lamports
+        let account_balance = 1000;
+        router.route_incoming_rewards(account_balance).unwrap();
+
+        // Verify rewards were routed correctly
+        assert_eq!(router.total_rewards(), 1000);
+        assert_eq!(router.reward_pool(), 1000);
+        assert_eq!(router.rewards_processed(), 0);
+
+        // Test routing additional 500 lamports
+        let account_balance = 1500;
+        router.route_incoming_rewards(account_balance).unwrap();
+
+        // Verify total rewards increased by difference
+        assert_eq!(router.total_rewards(), 1500);
+        assert_eq!(router.reward_pool(), 1500);
+        assert_eq!(router.rewards_processed(), 0);
+
+        // Test attempting to route with lower balance (should fail)
+        let result = router.route_incoming_rewards(1000);
+        assert!(result.is_err());
+
+        // Verify state didn't change after failed routing
+        assert_eq!(router.total_rewards(), 1500);
+        assert_eq!(router.reward_pool(), 1500);
+        assert_eq!(router.rewards_processed(), 0);
+    }
+}

core/src/ncn_fee_group.rs

@@ -16,7 +16,7 @@ pub enum NcnFeeGroupType {
     Reserved7 = 0x7,
 }
 
-#[derive(Debug, Clone, Copy, Zeroable, ShankType, Pod)]
+#[derive(Debug, Clone, Copy, Zeroable, ShankType, Pod, PartialEq, Eq)]
 #[repr(C)]
 pub struct NcnFeeGroup {
     pub group: u8,
@@ -221,24 +221,26 @@ mod tests {
             Err(TipRouterError::InvalidNcnFeeGroup)
         ));
     }
-
     #[test]
     fn test_all_groups() {
         let all_groups = NcnFeeGroup::all_groups();
 
         // Verify count matches FEE_GROUP_COUNT
         assert_eq!(all_groups.len(), NcnFeeGroup::FEE_GROUP_COUNT);
 
-        // Verify groups are in correct order
-        for (i, group) in all_groups.iter().enumerate() {
-            assert_eq!(group.group as usize, i);
-        }
+        // Verify groups are exactly as expected
+        let expected_groups = vec![
+            NcnFeeGroup::new(NcnFeeGroupType::Default),
+            NcnFeeGroup::new(NcnFeeGroupType::JTO),
+            NcnFeeGroup::new(NcnFeeGroupType::Reserved2),
+            NcnFeeGroup::new(NcnFeeGroupType::Reserved3),
+            NcnFeeGroup::new(NcnFeeGroupType::Reserved4),
+            NcnFeeGroup::new(NcnFeeGroupType::Reserved5),
+            NcnFeeGroup::new(NcnFeeGroupType::Reserved6),
+            NcnFeeGroup::new(NcnFeeGroupType::Reserved7),
+        ];
 
-        // Verify no duplicates
-        let mut unique_groups: Vec<u8> = all_groups.iter().map(|g| g.group).collect();
-        unique_groups.sort();
-        unique_groups.dedup();
-        assert_eq!(unique_groups.len(), NcnFeeGroup::FEE_GROUP_COUNT);
+        assert_eq!(all_groups, expected_groups);
     }
 
     #[test]

core/src/ncn_reward_router.rs

@@ -496,3 +496,54 @@ impl VaultRewardRoute {
         self.set_rewards(new_rewards)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use solana_program::pubkey::Pubkey;
+
+    use super::*;
+
+    #[test]
+    fn test_route_incoming_rewards() {
+        let mut router = NcnRewardRouter::new(
+            NcnFeeGroup::default(),
+            Pubkey::new_unique(), // ncn
+            Pubkey::new_unique(), // ncn
+            1,                    // ncn_epoch
+            1,                    // bump
+            100,                  // slot_created
+        );
+
+        // Initial state checks
+        assert_eq!(router.total_rewards(), 0);
+        assert_eq!(router.reward_pool(), 0);
+        assert_eq!(router.rewards_processed(), 0);
+
+        // Test routing 1000 lamports
+        let account_balance = 1000;
+        router.route_incoming_rewards(account_balance).unwrap();
+
+        // Verify rewards were routed correctly
+        assert_eq!(router.total_rewards(), 1000);
+        assert_eq!(router.reward_pool(), 1000);
+        assert_eq!(router.rewards_processed(), 0);
+
+        // Test routing additional 500 lamports
+        let account_balance = 1500;
+        router.route_incoming_rewards(account_balance).unwrap();
+
+        // Verify total rewards increased by difference
+        assert_eq!(router.total_rewards(), 1500);
+        assert_eq!(router.reward_pool(), 1500);
+        assert_eq!(router.rewards_processed(), 0);
+
+        // Test attempting to route with lower balance (should fail)
+        let result = router.route_incoming_rewards(1000);
+        assert!(result.is_err());
+
+        // Verify state didn't change after failed routing
+        assert_eq!(router.total_rewards(), 1500);
+        assert_eq!(router.reward_pool(), 1500);
+        assert_eq!(router.rewards_processed(), 0);
+    }
+}