@@ -574,10 +574,6 @@ mod tests {
574574
575575    type  E  = MainnetEthSpec ; 
576576
577-     // ============================================================================ 
578-     // Test Helpers and Utilities 
579-     // ============================================================================ 
580- 
581577    fn  setup_custody_context ( 
582578        spec :  & ChainSpec , 
583579        current_epoch :  Epoch , 
@@ -616,31 +612,6 @@ mod tests {
616612        } 
617613    } 
618614
619-     /// Helper function to test default CGC initialization for different node custody types. 
620- /// Verifies custody_group_count_at_head and num_of_custody_groups_to_sample match expectations. 
621- fn  assert_default_cgc_for_node_custody_type ( 
622-         node_custody_type :  NodeCustodyType , 
623-         expected_cgc_at_head :  u64 , 
624-         expected_sampling_size :  u64 , 
625-         spec :  & ChainSpec , 
626-     )  { 
627-         let  custody_context = CustodyContext :: < E > :: new ( node_custody_type,  spec) ; 
628- 
629-         assert_eq ! ( 
630-             custody_context. custody_group_count_at_head( spec) , 
631-             expected_cgc_at_head, 
632-             "custody_group_count_at_head should be {}" , 
633-             expected_cgc_at_head
634-         ) ; 
635- 
636-         assert_eq ! ( 
637-             custody_context. num_of_custody_groups_to_sample( Epoch :: new( 0 ) ,  spec) , 
638-             expected_sampling_size, 
639-             "num_of_custody_groups_to_sample should be {}" , 
640-             expected_sampling_size
641-         ) ; 
642-     } 
643- 
644615    /// Helper function to test CGC increases when switching node custody types. 
645616/// Verifies that CustodyCountChanged is returned with correct values and 
646617/// that custody_group_count_at_epoch returns appropriate values for current and next epoch. 
@@ -686,20 +657,15 @@ mod tests {
686657            "old_custody_group_count should be {}" , 
687658            persisted_cgc
688659        ) ; 
689- 
690-         // Verify effective epoch is next epoch 
691-         let  expected_effective_epoch = current_epoch + 1 ; 
692660        assert_eq ! ( 
693-             cgc_changed. effective_epoch,  expected_effective_epoch, 
661+             cgc_changed. effective_epoch, 
662+             current_epoch + 1 , 
694663            "effective epoch should be current_epoch + 1" 
695664        ) ; 
696- 
697-         // Verify sampling count 
698-         let  expected_sampling_count = spec
699-             . sampling_size_custody_groups ( expected_new_cgc) 
700-             . expect ( "should compute sampling size" ) ; 
701665        assert_eq ! ( 
702-             cgc_changed. sampling_count,  expected_sampling_count, 
666+             cgc_changed. sampling_count, 
667+             spec. sampling_size_custody_groups( expected_new_cgc) 
668+                 . expect( "should compute sampling size" ) , 
703669            "sampling_count should match expected value" 
704670        ) ; 
705671
@@ -755,47 +721,49 @@ mod tests {
755721        ) ; 
756722    } 
757723
758-     // ============================================================================ 
759-     // Default Initialization Tests 
760-     // ============================================================================ 
761- 
762724    #[ test]  
763725    fn  no_validators_supernode_default ( )  { 
764726        let  spec = E :: default_spec ( ) ; 
765-         assert_default_cgc_for_node_custody_type ( 
766-             NodeCustodyType :: Supernode , 
767-             spec. number_of_custody_groups , 
768-             spec. number_of_custody_groups , 
769-             & spec, 
727+         let  custody_context = CustodyContext :: < E > :: new ( NodeCustodyType :: Supernode ,  & spec) ; 
728+         assert_eq ! ( 
729+             custody_context. custody_group_count_at_head( & spec) , 
730+             spec. number_of_custody_groups
731+         ) ; 
732+         assert_eq ! ( 
733+             custody_context. num_of_custody_groups_to_sample( Epoch :: new( 0 ) ,  & spec) , 
734+             spec. number_of_custody_groups
770735        ) ; 
771736    } 
772737
773738    #[ test]  
774739    fn  no_validators_semi_supernode_default ( )  { 
775740        let  spec = E :: default_spec ( ) ; 
776-         assert_default_cgc_for_node_custody_type ( 
777-             NodeCustodyType :: SemiSupernode , 
778-             spec. number_of_custody_groups  / 2 , 
779-             spec. number_of_custody_groups  / 2 , 
780-             & spec, 
741+         let  custody_context = CustodyContext :: < E > :: new ( NodeCustodyType :: SemiSupernode ,  & spec) ; 
742+         assert_eq ! ( 
743+             custody_context. custody_group_count_at_head( & spec) , 
744+             spec. number_of_custody_groups / 2 
745+         ) ; 
746+         assert_eq ! ( 
747+             custody_context. num_of_custody_groups_to_sample( Epoch :: new( 0 ) ,  & spec) , 
748+             spec. number_of_custody_groups / 2 
781749        ) ; 
782750    } 
783751
784752    #[ test]  
785753    fn  no_validators_fullnode_default ( )  { 
786754        let  spec = E :: default_spec ( ) ; 
787-         assert_default_cgc_for_node_custody_type ( 
788-             NodeCustodyType :: Fullnode , 
755+         let  custody_context = CustodyContext :: < E > :: new ( NodeCustodyType :: Fullnode ,  & spec) ; 
756+         assert_eq ! ( 
757+             custody_context. custody_group_count_at_head( & spec) , 
789758            spec. custody_requirement, 
790-             spec. samples_per_slot , 
791-             & spec, 
759+             "head custody count should be minimum spec custody requirement" 
760+         ) ; 
761+         assert_eq ! ( 
762+             custody_context. num_of_custody_groups_to_sample( Epoch :: new( 0 ) ,  & spec) , 
763+             spec. samples_per_slot
792764        ) ; 
793765    } 
794766
795-     // ============================================================================ 
796-     // Validator Registration and CGC Updates 
797-     // ============================================================================ 
798- 
799767    #[ test]  
800768    fn  register_single_validator_should_update_cgc ( )  { 
801769        let  spec = E :: default_spec ( ) ; 
@@ -926,10 +894,6 @@ mod tests {
926894        ) ; 
927895    } 
928896
929-     // ============================================================================ 
930-     // Validator Expiry Tests 
931-     // ============================================================================ 
932- 
933897    #[ test]  
934898    fn  validator_dropped_after_no_registrations_within_expiry_should_not_reduce_cgc ( )  { 
935899        let  spec = E :: default_spec ( ) ; 
@@ -1022,10 +986,6 @@ mod tests {
1022986        ) ; 
1023987    } 
1024988
1025-     // ============================================================================ 
1026-     // Custody Columns and Data Column Initialization 
1027-     // ============================================================================ 
1028- 
1029989    #[ test]  
1030990    fn  should_init_ordered_data_columns_and_return_sampling_columns ( )  { 
1031991        let  spec = E :: default_spec ( ) ; 
@@ -1157,10 +1117,6 @@ mod tests {
11571117        ) ; 
11581118    } 
11591119
1160-     // ============================================================================ 
1161-     // Persistence and Restoration Tests 
1162-     // ============================================================================ 
1163- 
11641120    #[ test]  
11651121    fn  restore_from_persisted_fullnode_no_validators ( )  { 
11661122        let  spec = E :: default_spec ( ) ; 
@@ -1184,10 +1140,6 @@ mod tests {
11841140        ) ; 
11851141    } 
11861142
1187-     // ============================================================================ 
1188-     // Node Custody Type Switching Tests 
1189-     // ============================================================================ 
1190- 
11911143    /// Tests CLI flag change: Fullnode (CGC=0) → Supernode (CGC=128) 
11921144/// CGC should increase and trigger backfill via CustodyCountChanged. 
11931145#[ test]  
@@ -1340,10 +1292,6 @@ mod tests {
13401292        ) ; 
13411293    } 
13421294
1343-     // ============================================================================ 
1344-     // Historical Custody and Backfill Tests 
1345-     // ============================================================================ 
1346- 
13471295    #[ test]  
13481296    fn  restore_with_validator_custody_history_across_epochs ( )  { 
13491297        let  spec = E :: default_spec ( ) ; 
0 commit comments