test new signature scheme in contract unit tests

Author: kevindeforth

crates/contract/src/lib.rs

@@ -2393,10 +2393,11 @@ mod tests {
             }
         }
     }
-
+    const NUM_GENERATED_DOMAINS: usize = 1;
+    const NUM_DOMAINS: usize = 5;
     #[test]
     fn test_start_node_migration_failure_not_participant() {
-        let running_state = ProtocolContractState::Running(gen_running_state(2));
+        let running_state = ProtocolContractState::Running(gen_running_state(NUM_DOMAINS));
         let mut contract = MpcContract::new_from_protocol_state(running_state);
 
         // sanity check
@@ -2413,7 +2414,7 @@ mod tests {
 
     #[test]
     fn test_start_node_migration_success() {
-        let running_state = ProtocolContractState::Running(gen_running_state(2));
+        let running_state = ProtocolContractState::Running(gen_running_state(NUM_DOMAINS));
         let mut contract = MpcContract::new_from_protocol_state(running_state);
 
         // sanity check
@@ -2469,14 +2470,14 @@ mod tests {
     #[test]
     fn test_start_node_migration_failure_initializing() {
         let initializing_state =
-            ProtocolContractState::Initializing(gen_initializing_state(2, 0).1);
+            ProtocolContractState::Initializing(gen_initializing_state(NUM_DOMAINS, 0).1);
         let contract = MpcContract::new_from_protocol_state(initializing_state);
         test_start_migration_node_failure_not_running(contract);
     }
 
     #[test]
     fn test_start_node_migration_failure_resharing() {
-        let resharing_state = ProtocolContractState::Resharing(gen_resharing_state(2).1);
+        let resharing_state = ProtocolContractState::Resharing(gen_resharing_state(NUM_DOMAINS).1);
         let contract = MpcContract::new_from_protocol_state(resharing_state);
         test_start_migration_node_failure_not_running(contract);
     }
@@ -2498,22 +2499,22 @@ mod tests {
 
     #[test]
     fn test_register_backup_service_fail_non_participant_running() {
-        let running_state = ProtocolContractState::Running(gen_running_state(2));
+        let running_state = ProtocolContractState::Running(gen_running_state(NUM_DOMAINS));
         let contract = MpcContract::new_from_protocol_state(running_state);
         test_register_backup_service_fail_non_participant(contract);
     }
 
     #[test]
     fn test_register_backup_service_fail_non_participant_initializing() {
         let initializing_state =
-            ProtocolContractState::Initializing(gen_initializing_state(2, 0).1);
+            ProtocolContractState::Initializing(gen_initializing_state(NUM_DOMAINS, 0).1);
         let contract = MpcContract::new_from_protocol_state(initializing_state);
         test_register_backup_service_fail_non_participant(contract);
     }
 
     #[test]
     fn test_register_backup_service_fail_non_participant_resharnig() {
-        let resharing_state = ProtocolContractState::Resharing(gen_resharing_state(2).1);
+        let resharing_state = ProtocolContractState::Resharing(gen_resharing_state(NUM_DOMAINS).1);
         let contract = MpcContract::new_from_protocol_state(resharing_state);
         test_register_backup_service_fail_non_participant(contract);
     }
@@ -2549,7 +2550,7 @@ mod tests {
 
     #[test]
     fn test_register_backup_service_success_running() {
-        let running_state = gen_running_state(2);
+        let running_state = gen_running_state(NUM_DOMAINS);
         let participants = running_state.parameters.participants().clone();
         let running_state = ProtocolContractState::Running(running_state);
         let contract = MpcContract::new_from_protocol_state(running_state);
@@ -2558,7 +2559,7 @@ mod tests {
 
     #[test]
     fn test_register_backup_service_success_resharing() {
-        let resharing_state = gen_resharing_state(2).1;
+        let resharing_state = gen_resharing_state(NUM_DOMAINS).1;
         let participants = resharing_state
             .resharing_key
             .proposed_parameters()
@@ -2571,7 +2572,7 @@ mod tests {
 
     #[test]
     fn test_register_backup_service_success_initializing() {
-        let initializing_state = gen_initializing_state(2, 0).1;
+        let initializing_state = gen_initializing_state(NUM_DOMAINS, 0).1;
         let participants = initializing_state
             .generating_key
             .proposed_parameters()
@@ -2584,7 +2585,7 @@ mod tests {
 
     #[test]
     fn test_conclude_node_migration_success() {
-        let running_state = gen_running_state(2);
+        let running_state = gen_running_state(NUM_DOMAINS);
         let keyset = running_state.keyset.clone();
         let participants = running_state.parameters.participants().clone();
         let running_state = ProtocolContractState::Running(running_state);
@@ -2614,7 +2615,7 @@ mod tests {
 
     #[test]
     fn test_conclude_node_migration_invalid_tee() {
-        let running_state = gen_running_state(2);
+        let running_state = gen_running_state(NUM_DOMAINS);
         let keyset = running_state.keyset.clone();
         let participants = running_state.parameters.participants().clone();
         let running_state = ProtocolContractState::Running(running_state);
@@ -2642,7 +2643,7 @@ mod tests {
 
     #[test]
     fn test_conclude_node_migration_migration_not_found() {
-        let running_state = gen_running_state(2);
+        let running_state = gen_running_state(NUM_DOMAINS);
         let keyset = running_state.keyset.clone();
         let participants = running_state.parameters.participants().clone();
         let running_state = ProtocolContractState::Running(running_state);
@@ -2675,7 +2676,7 @@ mod tests {
 
     #[test]
     fn test_conclude_node_migration_keyset_mismatch() {
-        let running_state = gen_running_state(2);
+        let running_state = gen_running_state(NUM_DOMAINS);
         let mut keyset = running_state.keyset.clone();
         keyset.epoch_id = keyset.epoch_id.next();
         let participants = running_state.parameters.participants().clone();
@@ -2709,7 +2710,7 @@ mod tests {
 
     #[test]
     fn test_conclude_node_migration_not_participant() {
-        let running_state = gen_running_state(2);
+        let running_state = gen_running_state(NUM_DOMAINS);
         let keyset = running_state.keyset.clone();
         let running_state = ProtocolContractState::Running(running_state);
         let mut contract = MpcContract::new_from_protocol_state(running_state);
@@ -2764,7 +2765,7 @@ mod tests {
 
     #[test]
     fn test_conclude_node_migration_failure_resharing() {
-        let (_, resharing_state) = gen_resharing_state(2);
+        let (_, resharing_state) = gen_resharing_state(NUM_DOMAINS);
 
         let keyset = resharing_state.previous_running_state.keyset.clone();
         let participants = resharing_state
@@ -2779,7 +2780,7 @@ mod tests {
 
     #[test]
     fn test_conclude_node_migration_failure_initializing() {
-        let (_, initializing) = gen_initializing_state(2, 0);
+        let (_, initializing) = gen_initializing_state(NUM_DOMAINS, 0);
 
         let keyset = Keyset::new(
             initializing.generating_key.epoch_id(),
@@ -2868,7 +2869,7 @@ mod tests {
 
     #[test]
     pub fn test_cleanup_orphaned_node_migrations() {
-        let running_state = gen_running_state(2);
+        let running_state = gen_running_state(NUM_DOMAINS);
         let participants = running_state.parameters.participants().clone();
         let running_state = ProtocolContractState::Running(running_state);
         let mut contract = MpcContract::new_from_protocol_state(running_state);
@@ -3001,26 +3002,29 @@ mod tests {
 
     #[test]
     pub fn test_proposed_updates_interface_running() {
-        let protocol_contract_state = ProtocolContractState::Running(gen_running_state(2));
+        let protocol_contract_state =
+            ProtocolContractState::Running(gen_running_state(NUM_DOMAINS));
         test_proposed_updates_case_given_state(protocol_contract_state);
     }
 
     #[test]
     pub fn test_proposed_updates_interface_resharing() {
-        let protocol_contract_state = ProtocolContractState::Resharing(gen_resharing_state(2).1);
+        let protocol_contract_state =
+            ProtocolContractState::Resharing(gen_resharing_state(NUM_DOMAINS).1);
         test_proposed_updates_case_given_state(protocol_contract_state);
     }
 
     #[test]
     pub fn test_proposed_updates_interface_initialzing() {
-        let protocol_contract_state =
-            ProtocolContractState::Initializing(gen_initializing_state(2, 1).1);
+        let protocol_contract_state = ProtocolContractState::Initializing(
+            gen_initializing_state(NUM_DOMAINS, NUM_GENERATED_DOMAINS).1,
+        );
         test_proposed_updates_case_given_state(protocol_contract_state);
     }
 
     #[test]
     pub fn test_remove_update_vote_running() {
-        let running_state = gen_running_state(2);
+        let running_state = gen_running_state(NUM_DOMAINS);
         let participants = running_state.parameters.participants().clone();
         let protocol_contract_state = ProtocolContractState::Running(running_state);
         let mut contract = MpcContract::new_from_protocol_state(protocol_contract_state);
@@ -3061,7 +3065,7 @@ mod tests {
     #[test]
     #[should_panic(expected = "not a voter")]
     fn test_remove_update_vote_panics_if_non_voter() {
-        let running_state = gen_running_state(2);
+        let running_state = gen_running_state(NUM_DOMAINS);
         let protocol_contract_state = ProtocolContractState::Running(running_state);
         let mut contract = MpcContract::new_from_protocol_state(protocol_contract_state);
         let expected = propose_and_vote_code(0, &mut contract);
@@ -3080,7 +3084,8 @@ mod tests {
     #[test]
     #[should_panic(expected = "protocol must be in running state")]
     pub fn test_remove_update_vote_resharing() {
-        let protocol_contract_state = ProtocolContractState::Resharing(gen_resharing_state(2).1);
+        let protocol_contract_state =
+            ProtocolContractState::Resharing(gen_resharing_state(NUM_DOMAINS).1);
         let mut contract = MpcContract::new_from_protocol_state(protocol_contract_state);
         let account_id = gen_account_id();
         testing_env!(VMContextBuilder::new()
@@ -3093,8 +3098,9 @@ mod tests {
     #[test]
     #[should_panic(expected = "protocol must be in running state")]
     pub fn test_remove_update_vote_initializing() {
-        let protocol_contract_state =
-            ProtocolContractState::Initializing(gen_initializing_state(2, 1).1);
+        let protocol_contract_state = ProtocolContractState::Initializing(
+            gen_initializing_state(NUM_DOMAINS, NUM_GENERATED_DOMAINS).1,
+        );
         let mut contract = MpcContract::new_from_protocol_state(protocol_contract_state);
         let account_id = gen_account_id();
         testing_env!(VMContextBuilder::new()
@@ -3164,7 +3170,7 @@ mod tests {
     /// (simulating post-resharing cleanup).
     #[test]
     pub fn test_remove_non_participant_update_votes() {
-        let running_state = gen_running_state(2);
+        let running_state = gen_running_state(NUM_DOMAINS);
         let participants = running_state.parameters.participants().clone();
         let mut contract =
             MpcContract::new_from_protocol_state(ProtocolContractState::Running(running_state));
@@ -3216,9 +3222,9 @@ mod tests {
     }
 
     #[rstest]
-    #[case(ProtocolContractState::Running(gen_running_state(2)))]
-    #[case(ProtocolContractState::Resharing(gen_resharing_state(2).1))]
-    #[case(ProtocolContractState::Initializing(gen_initializing_state(2, 1).1))]
+    #[case(ProtocolContractState::Running(gen_running_state(NUM_DOMAINS)))]
+    #[case(ProtocolContractState::Resharing(gen_resharing_state(NUM_DOMAINS).1))]
+    #[case(ProtocolContractState::Initializing(gen_initializing_state(NUM_DOMAINS, NUM_GENERATED_DOMAINS).1))]
     fn test_contract_stores_allowed_hashes(#[case] protocol_state: ProtocolContractState) {
         const CURRENT_BLOCK_TIME_STAMP: u64 = 10;
         let mut contract = MpcContract::new_from_protocol_state(protocol_state);

crates/contract/src/primitives/domain.rs

@@ -220,11 +220,11 @@ pub mod tests {
         let domains2 = vec![
             DomainConfig {
                 id: DomainId(2),
-                scheme: SignatureScheme::Secp256k1,
+                scheme: SignatureScheme::Bls12381,
             },
             DomainConfig {
                 id: DomainId(3),
-                scheme: SignatureScheme::Ed25519,
+                scheme: SignatureScheme::V2Secp256k1,
             },
         ];
         let new_registry = new_registry.add_domains(domains2.clone()).unwrap();
@@ -254,57 +254,30 @@ pub mod tests {
 
     #[test]
     fn test_retain_domains() {
-        let mut registry = DomainRegistry::from_raw_validated(
-            vec![
-                DomainConfig {
-                    id: DomainId(0),
-                    scheme: SignatureScheme::Secp256k1,
-                },
-                DomainConfig {
-                    id: DomainId(2),
-                    scheme: SignatureScheme::Ed25519,
-                },
-                DomainConfig {
-                    id: DomainId(3),
-                    scheme: SignatureScheme::Secp256k1,
-                },
-            ],
-            6,
-        )
-        .unwrap();
+        let expected = vec![
+            DomainConfig {
+                id: DomainId(0),
+                scheme: SignatureScheme::Secp256k1,
+            },
+            DomainConfig {
+                id: DomainId(2),
+                scheme: SignatureScheme::Ed25519,
+            },
+            DomainConfig {
+                id: DomainId(3),
+                scheme: SignatureScheme::Bls12381,
+            },
+            DomainConfig {
+                id: DomainId(4),
+                scheme: SignatureScheme::V2Secp256k1,
+            },
+        ];
+        let mut registry = DomainRegistry::from_raw_validated(expected.clone(), 6).unwrap();
         registry.retain_domains(3);
-        assert_eq!(
-            registry.domains,
-            vec![
-                DomainConfig {
-                    id: DomainId(0),
-                    scheme: SignatureScheme::Secp256k1,
-                },
-                DomainConfig {
-                    id: DomainId(2),
-                    scheme: SignatureScheme::Ed25519,
-                },
-                DomainConfig {
-                    id: DomainId(3),
-                    scheme: SignatureScheme::Secp256k1,
-                },
-            ]
-        );
+        assert_eq!(registry.domains, expected[0..3]);
         assert_eq!(registry.next_domain_id, 6);
         registry.retain_domains(2);
-        assert_eq!(
-            registry.domains,
-            vec![
-                DomainConfig {
-                    id: DomainId(0),
-                    scheme: SignatureScheme::Secp256k1,
-                },
-                DomainConfig {
-                    id: DomainId(2),
-                    scheme: SignatureScheme::Ed25519,
-                },
-            ]
-        );
+        assert_eq!(registry.domains, expected[0..2]);
         assert_eq!(registry.next_domain_id, 6);
         registry.retain_domains(0);
         assert_eq!(registry.domains, Vec::new());

crates/contract/src/primitives/test_utils.rs

@@ -13,10 +13,11 @@ use std::collections::BTreeMap;
 
 use super::domain::{DomainConfig, DomainId, DomainRegistry, SignatureScheme};
 
-const ALL_PROTOCOLS: [SignatureScheme; 3] = [
+const ALL_PROTOCOLS: [SignatureScheme; 4] = [
     SignatureScheme::Secp256k1,
     SignatureScheme::Ed25519,
     SignatureScheme::Bls12381,
+    SignatureScheme::V2Secp256k1,
 ];
 
 /// Generates a valid DomainRegistry with various signature schemes, with num_domains total.

crates/contract/src/state.rs

@@ -302,6 +302,26 @@ impl ProtocolContractState {
             _ => None,
         }
     }
+
+    pub fn next_domain_id(&self) -> u64 {
+        match self {
+            ProtocolContractState::Running(state) => state.domains.next_domain_id(),
+            ProtocolContractState::Initializing(state) => state.domains.next_domain_id(),
+            ProtocolContractState::Resharing(state) => {
+                state.previous_running_state.domains.next_domain_id()
+            }
+            _ => panic!("should be in initializing state"),
+        }
+    }
+
+    pub fn current_epoch(&self) -> EpochId {
+        match self {
+            ProtocolContractState::Running(state) => state.keyset.epoch_id,
+            ProtocolContractState::Initializing(state) => state.epoch_id,
+            ProtocolContractState::Resharing(state) => state.previous_running_state.keyset.epoch_id,
+            _ => panic!("should be in initializing state"),
+        }
+    }
 }
 
 #[derive(Debug, Clone)]

crates/contract/src/state/initializing.rs

@@ -336,9 +336,10 @@ pub mod tests {
         test_initializing_contract_state_for(3, 2);
     }
 
+    const NUM_DOMAINS: usize = 5;
     #[test]
     fn test_cancel_key_generation() {
-        let (mut env, mut state) = gen_initializing_state(5, 2);
+        let (mut env, mut state) = gen_initializing_state(NUM_DOMAINS, 2);
 
         // Vote for domain #2.
         let leader = find_leader(&state.generating_key);