From dd51d0a2f226dbb79d217f973e933c106244f1a4 Mon Sep 17 00:00:00 2001
From: Matan Markind <matan@starkware.co>
Date: Wed, 19 Jun 2024 13:51:52 +0300
Subject: [PATCH] refactor(consensus): add test util to more cleanly create
 validator ids

---
 .../src/single_height_consensus_test.rs            | 14 +++++++-------
 .../sequencing/papyrus_consensus/src/test_utils.rs |  7 +++++++
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/crates/sequencing/papyrus_consensus/src/single_height_consensus_test.rs b/crates/sequencing/papyrus_consensus/src/single_height_consensus_test.rs
index d1c0fa5c9a..cfeb99bd81 100644
--- a/crates/sequencing/papyrus_consensus/src/single_height_consensus_test.rs
+++ b/crates/sequencing/papyrus_consensus/src/single_height_consensus_test.rs
@@ -6,19 +6,19 @@ use starknet_types_core::felt::Felt;
 use tokio;
 
 use super::SingleHeightConsensus;
-use crate::test_utils::{MockTestContext, TestBlock};
-use crate::types::{ConsensusBlock, ProposalInit, ValidatorId};
+use crate::test_utils::{validator_id, MockTestContext, TestBlock};
+use crate::types::{ConsensusBlock, ProposalInit};
 
 #[tokio::test]
 async fn proposer() {
     let mut context = MockTestContext::new();
 
-    let node_id: ValidatorId = 1_u32.into();
+    let node_id = validator_id(1);
     let block = TestBlock { content: vec![1, 2, 3], id: BlockHash(Felt::ONE) };
     // Set expectations for how the test should run:
     context
         .expect_validators()
-        .returning(move |_| vec![node_id, 2_u32.into(), 3_u32.into(), 4_u32.into()]);
+        .returning(move |_| vec![node_id, validator_id(2), validator_id(3), validator_id(4)]);
     context.expect_proposer().returning(move |_, _| node_id);
     let block_clone = block.clone();
     context.expect_build_proposal().returning(move |_| {
@@ -52,14 +52,14 @@ async fn proposer() {
 async fn validator() {
     let mut context = MockTestContext::new();
 
-    let node_id: ValidatorId = 1_u32.into();
-    let proposer: ValidatorId = 2_u32.into();
+    let node_id = validator_id(1);
+    let proposer = validator_id(2);
     let block = TestBlock { content: vec![1, 2, 3], id: BlockHash(Felt::ONE) };
 
     // Set expectations for how the test should run:
     context
         .expect_validators()
-        .returning(move |_| vec![node_id, proposer, 3_u32.into(), 4_u32.into()]);
+        .returning(move |_| vec![node_id, proposer, validator_id(3), validator_id(4)]);
     context.expect_proposer().returning(move |_, _| proposer);
     let block_clone = block.clone();
     context.expect_validate_proposal().returning(move |_, _| {
diff --git a/crates/sequencing/papyrus_consensus/src/test_utils.rs b/crates/sequencing/papyrus_consensus/src/test_utils.rs
index a86587ea4c..e2f1eaeeca 100644
--- a/crates/sequencing/papyrus_consensus/src/test_utils.rs
+++ b/crates/sequencing/papyrus_consensus/src/test_utils.rs
@@ -2,6 +2,7 @@ use async_trait::async_trait;
 use futures::channel::{mpsc, oneshot};
 use mockall::mock;
 use starknet_api::block::{BlockHash, BlockNumber};
+use starknet_api::core::ContractAddress;
 
 use crate::types::{ConsensusBlock, ConsensusContext, ConsensusError, ProposalInit, ValidatorId};
 
@@ -56,3 +57,9 @@ mock! {
         ) -> Result<(), ConsensusError>;
     }
 }
+
+// Convenience function for creating a `ValidatorId` since i32 is the default type for explicit
+// numbers.
+pub fn validator_id(id: u32) -> ValidatorId {
+    ContractAddress::from(id)
+}