Add tests

tests/src/mock_validator.rs

@@ -36,11 +36,16 @@ impl MockValidator {
     }
 
     pub async fn do_register_validator(&self) -> Result<(), Error> {
-        let url = self.comm_boost.register_validator_url().unwrap();
+        self.do_register_custom_validators(vec![]).await
+    }
 
-        let registration: Vec<ValidatorRegistration> = vec![];
+    pub async fn do_register_custom_validators(
+        &self,
+        registrations: Vec<ValidatorRegistration>,
+    ) -> Result<(), Error> {
+        let url = self.comm_boost.register_validator_url().unwrap();
 
-        self.comm_boost.client.post(url).json(&registration).send().await?.error_for_status()?;
+        self.comm_boost.client.post(url).json(&registrations).send().await?.error_for_status()?;
 
         Ok(())
     }

tests/src/utils.rs

@@ -32,3 +32,22 @@ pub fn generate_mock_relay(port: u16, pubkey: BlsPublicKey) -> Result<RelayClien
     };
     RelayClient::new(config)
 }
+
+pub fn generate_mock_relay_with_batch_size(
+    port: u16,
+    pubkey: BlsPublicKey,
+    batch_size: usize,
+) -> Result<RelayClient> {
+    let entry =
+        RelayEntry { id: format!("mock_{port}"), pubkey, url: get_local_address(port).parse()? };
+    let config = RelayConfig {
+        entry,
+        id: None,
+        headers: None,
+        enable_timing_games: false,
+        target_first_request_ms: None,
+        frequency_get_header_ms: None,
+        validator_registration_batch_size: Some(batch_size),
+    };
+    RelayClient::new(config)
+}

tests/tests/pbs_integration.rs

@@ -6,7 +6,7 @@ use std::{
     u64,
 };
 
-use alloy::primitives::U256;
+use alloy::{primitives::U256, rpc::types::beacon::relay::ValidatorRegistration};
 use cb_common::{
     config::{PbsConfig, PbsModuleConfig, RuntimeMuxConfig},
     pbs::RelayClient,
@@ -18,7 +18,7 @@ use cb_pbs::{DefaultBuilderApi, PbsService, PbsState};
 use cb_tests::{
     mock_relay::{start_mock_relay_service, MockRelayState},
     mock_validator::MockValidator,
-    utils::{generate_mock_relay, setup_test_env},
+    utils::{generate_mock_relay, generate_mock_relay_with_batch_size, setup_test_env},
 };
 use eyre::Result;
 use tracing::info;
@@ -143,6 +143,48 @@ async fn test_register_validators() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_batch_register_validators() -> Result<()> {
+    setup_test_env();
+    let signer = random_secret();
+    let pubkey: BlsPublicKey = blst_pubkey_to_alloy(&signer.sk_to_pk()).into();
+
+    let chain = Chain::Holesky;
+    let port = 3310;
+
+    let relays = vec![generate_mock_relay_with_batch_size(port + 1, *pubkey, 5)?];
+    let mock_state = Arc::new(MockRelayState::new(chain, signer));
+    tokio::spawn(start_mock_relay_service(mock_state.clone(), port + 1));
+
+    let config = to_pbs_config(chain, get_pbs_static_config(port), relays);
+    let state = PbsState::new(config);
+    tokio::spawn(PbsService::run::<(), DefaultBuilderApi>(state));
+
+    // leave some time to start servers
+    tokio::time::sleep(Duration::from_millis(100)).await;
+
+    let data = include_str!("../data/registration_holesky.json");
+    let registrations: Vec<ValidatorRegistration> = serde_json::from_str(data)?;
+
+    let mock_validator = MockValidator::new(port)?;
+    info!("Sending register validator");
+    let res = mock_validator.do_register_custom_validators(registrations.clone()).await;
+
+    // registrations.len() == 17. 5 per batch, 4 batches
+    assert!(res.is_ok());
+    assert_eq!(mock_state.received_register_validator(), 4);
+
+    let mock_validator = MockValidator::new(port)?;
+    info!("Sending register validator");
+    let res = mock_validator.do_register_custom_validators(registrations[..2].to_vec()).await;
+
+    // Expected one more registration request
+    assert!(res.is_ok());
+    assert_eq!(mock_state.received_register_validator(), 5);
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn test_submit_block() -> Result<()> {
     setup_test_env();