approval-voting improvement: include all tranche0 assignments in one certificate (#1178)

**_PR migrated from https://github.com/paritytech/polkadot/pull/6782_** 

This PR will upgrade the network protocol to version 3 -> VStaging which
will later be renamed to V3. This version introduces a new kind of
assignment certificate that will be used for tranche0 assignments.
Instead of issuing/importing one tranche0 assignment per candidate,
there will be just one certificate per relay chain block per validator.
However, we will not be sending out the new assignment certificates,
yet. So everything should work exactly as before. Once the majority of
the validators have been upgraded to the new protocol version we will
enable the new certificates (starting at a specific relay chain block)
with a new client update.

There are still a few things that need to be done:

- [x] Use bitfield instead of Vec<CandidateIndex>:
https://github.com/paritytech/polkadot/pull/6802
  - [x] Fix existing approval-distribution and approval-voting tests
  - [x] Fix bitfield-distribution and statement-distribution tests
  - [x] Fix network bridge tests
  - [x] Implement todos in the code
  - [x] Add tests to cover new code
  - [x] Update metrics
  - [x] Remove the approval distribution aggression levels: TBD PR
  - [x] Parachains DB migration 
  - [x] Test network protocol upgrade on Versi
  - [x] Versi Load test
  - [x] Add Zombienet test
  - [x] Documentation updates
- [x] Fix for sending DistributeAssignment for each candidate claimed by
a v2 assignment (warning: Importing locally an already known assignment)
 - [x]  Fix AcceptedDuplicate
 - [x] Fix DB migration so that we can still keep old data.
 - [x] Final Versi burn in

---------

Signed-off-by: Andrei Sandu <andrei-mihail@parity.io>
Signed-off-by: Alexandru Gheorghe <alexandru.gheorghe@parity.io>
Co-authored-by: Alexandru Gheorghe <alexandru.gheorghe@parity.io>

polkadot/node/network/approval-distribution/src/metrics.rs

@@ -15,6 +15,7 @@
 // along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
 
 use polkadot_node_metrics::metrics::{prometheus, Metrics as MetricsTrait};
+use polkadot_node_primitives::approval::v2::AssignmentCertKindV2;
 
 /// Approval Distribution metrics.
 #[derive(Default, Clone)]
@@ -22,21 +23,34 @@ pub struct Metrics(Option<MetricsInner>);
 
 #[derive(Clone)]
 struct MetricsInner {
-	assignments_imported_total: prometheus::Counter<prometheus::U64>,
+	assignments_imported_total: prometheus::CounterVec<prometheus::U64>,
 	approvals_imported_total: prometheus::Counter<prometheus::U64>,
 	unified_with_peer_total: prometheus::Counter<prometheus::U64>,
 	aggression_l1_messages_total: prometheus::Counter<prometheus::U64>,
 	aggression_l2_messages_total: prometheus::Counter<prometheus::U64>,
-
 	time_unify_with_peer: prometheus::Histogram,
 	time_import_pending_now_known: prometheus::Histogram,
 	time_awaiting_approval_voting: prometheus::Histogram,
 }
 
+trait AsLabel {
+	fn as_label(&self) -> &str;
+}
+
+impl AsLabel for &AssignmentCertKindV2 {
+	fn as_label(&self) -> &str {
+		match self {
+			AssignmentCertKindV2::RelayVRFDelay { .. } => "VRF Delay",
+			AssignmentCertKindV2::RelayVRFModulo { .. } => "VRF Modulo",
+			AssignmentCertKindV2::RelayVRFModuloCompact { .. } => "VRF Modulo Compact",
+		}
+	}
+}
+
 impl Metrics {
-	pub(crate) fn on_assignment_imported(&self) {
+	pub(crate) fn on_assignment_imported(&self, kind: &AssignmentCertKindV2) {
 		if let Some(metrics) = &self.0 {
-			metrics.assignments_imported_total.inc();
+			metrics.assignments_imported_total.with_label_values(&[kind.as_label()]).inc();
 		}
 	}
 
@@ -89,9 +103,12 @@ impl MetricsTrait for Metrics {
 	fn try_register(registry: &prometheus::Registry) -> Result<Self, prometheus::PrometheusError> {
 		let metrics = MetricsInner {
 			assignments_imported_total: prometheus::register(
-				prometheus::Counter::new(
-					"polkadot_parachain_assignments_imported_total",
-					"Number of valid assignments imported locally or from other peers.",
+				prometheus::CounterVec::new(
+					prometheus::Opts::new(
+						"polkadot_parachain_assignments_imported_total",
+						"Number of valid assignments imported locally or from other peers.",
+					),
+					&["kind"],
 				)?,
 				registry,
 			)?,
@@ -124,10 +141,16 @@ impl MetricsTrait for Metrics {
 				registry,
 			)?,
 			time_unify_with_peer: prometheus::register(
-				prometheus::Histogram::with_opts(prometheus::HistogramOpts::new(
-					"polkadot_parachain_time_unify_with_peer",
-					"Time spent within fn `unify_with_peer`.",
-				).buckets(vec![0.000625, 0.00125,0.0025, 0.005, 0.0075, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0,]))?,
+				prometheus::Histogram::with_opts(
+					prometheus::HistogramOpts::new(
+						"polkadot_parachain_time_unify_with_peer",
+						"Time spent within fn `unify_with_peer`.",
+					)
+					.buckets(vec![
+						0.000625, 0.00125, 0.0025, 0.005, 0.0075, 0.01, 0.025, 0.05, 0.1, 0.25,
+						0.5, 1.0, 2.5, 5.0, 10.0,
+					]),
+				)?,
 				registry,
 			)?,
 			time_import_pending_now_known: prometheus::register(

polkadot/node/network/approval-distribution/src/tests.rs

@@ -24,20 +24,26 @@ use polkadot_node_network_protocol::{
 	view, ObservedRole,
 };
 use polkadot_node_primitives::approval::{
-	AssignmentCertKind, VrfOutput, VrfProof, VrfSignature, RELAY_VRF_MODULO_CONTEXT,
+	v1::{
+		AssignmentCert, AssignmentCertKind, IndirectAssignmentCert, VrfOutput, VrfProof,
+		VrfSignature,
+	},
+	v2::{
+		AssignmentCertKindV2, AssignmentCertV2, CoreBitfield, IndirectAssignmentCertV2,
+		RELAY_VRF_MODULO_CONTEXT,
+	},
 };
 use polkadot_node_subsystem::messages::{
 	network_bridge_event, AllMessages, ApprovalCheckError, ReportPeerMessage,
 };
 use polkadot_node_subsystem_test_helpers as test_helpers;
 use polkadot_node_subsystem_util::{reputation::add_reputation, TimeoutExt as _};
-use polkadot_primitives::{AuthorityDiscoveryId, BlakeTwo256, HashT};
+use polkadot_primitives::{AuthorityDiscoveryId, BlakeTwo256, CoreIndex, HashT};
 use polkadot_primitives_test_helpers::dummy_signature;
 use rand::SeedableRng;
 use sp_authority_discovery::AuthorityPair as AuthorityDiscoveryPair;
 use sp_core::crypto::Pair as PairT;
 use std::time::Duration;
-
 type VirtualOverseer = test_helpers::TestSubsystemContextHandle<ApprovalDistributionMessage>;
 
 fn test_harness<T: Future<Output = VirtualOverseer>>(
@@ -219,15 +225,15 @@ async fn setup_gossip_topology(
 async fn setup_peer_with_view(
 	virtual_overseer: &mut VirtualOverseer,
 	peer_id: &PeerId,
-	validation_version: ValidationVersion,
 	view: View,
+	version: ValidationVersion,
 ) {
 	overseer_send(
 		virtual_overseer,
 		ApprovalDistributionMessage::NetworkBridgeUpdate(NetworkBridgeEvent::PeerConnected(
 			*peer_id,
 			ObservedRole::Full,
-			validation_version.into(),
+			version.into(),
 			None,
 		)),
 	)
@@ -244,12 +250,43 @@ async fn setup_peer_with_view(
 async fn send_message_from_peer(
 	virtual_overseer: &mut VirtualOverseer,
 	peer_id: &PeerId,
-	msg: net_protocol::ApprovalDistributionMessage,
+	msg: protocol_v1::ApprovalDistributionMessage,
 ) {
 	overseer_send(
 		virtual_overseer,
 		ApprovalDistributionMessage::NetworkBridgeUpdate(NetworkBridgeEvent::PeerMessage(
-			*peer_id, msg,
+			*peer_id,
+			Versioned::V1(msg),
+		)),
+	)
+	.await;
+}
+
+async fn send_message_from_peer_v2(
+	virtual_overseer: &mut VirtualOverseer,
+	peer_id: &PeerId,
+	msg: protocol_v2::ApprovalDistributionMessage,
+) {
+	overseer_send(
+		virtual_overseer,
+		ApprovalDistributionMessage::NetworkBridgeUpdate(NetworkBridgeEvent::PeerMessage(
+			*peer_id,
+			Versioned::V2(msg),
+		)),
+	)
+	.await;
+}
+
+async fn send_message_from_peer_vstaging(
+	virtual_overseer: &mut VirtualOverseer,
+	peer_id: &PeerId,
+	msg: protocol_vstaging::ApprovalDistributionMessage,
+) {
+	overseer_send(
+		virtual_overseer,
+		ApprovalDistributionMessage::NetworkBridgeUpdate(NetworkBridgeEvent::PeerMessage(
+			*peer_id,
+			Versioned::VStaging(msg),
 		)),
 	)
 	.await;
@@ -273,6 +310,28 @@ fn fake_assignment_cert(block_hash: Hash, validator: ValidatorIndex) -> Indirect
 	}
 }
 
+fn fake_assignment_cert_v2(
+	block_hash: Hash,
+	validator: ValidatorIndex,
+	core_bitfield: CoreBitfield,
+) -> IndirectAssignmentCertV2 {
+	let ctx = schnorrkel::signing_context(RELAY_VRF_MODULO_CONTEXT);
+	let msg = b"WhenParachains?";
+	let mut prng = rand_core::OsRng;
+	let keypair = schnorrkel::Keypair::generate_with(&mut prng);
+	let (inout, proof, _) = keypair.vrf_sign(ctx.bytes(msg));
+	let out = inout.to_output();
+
+	IndirectAssignmentCertV2 {
+		block_hash,
+		validator,
+		cert: AssignmentCertV2 {
+			kind: AssignmentCertKindV2::RelayVRFModuloCompact { core_bitfield },
+			vrf: VrfSignature { output: VrfOutput(out), proof: VrfProof(proof) },
+		},
+	}
+}
+
 async fn expect_reputation_change(
 	virtual_overseer: &mut VirtualOverseer,
 	peer_id: &PeerId,
@@ -331,9 +390,9 @@ fn try_import_the_same_assignment() {
 	let _ = test_harness(state_without_reputation_delay(), |mut virtual_overseer| async move {
 		let overseer = &mut virtual_overseer;
 		// setup peers
-		setup_peer_with_view(overseer, &peer_a, ValidationVersion::V1, view![]).await;
-		setup_peer_with_view(overseer, &peer_b, ValidationVersion::V1, view![hash]).await;
-		setup_peer_with_view(overseer, &peer_c, ValidationVersion::V1, view![hash]).await;
+		setup_peer_with_view(overseer, &peer_a, view![], ValidationVersion::V1).await;
+		setup_peer_with_view(overseer, &peer_b, view![hash], ValidationVersion::V1).await;
+		setup_peer_with_view(overseer, &peer_c, view![hash], ValidationVersion::V1).await;
 
 		// new block `hash_a` with 1 candidates
 		let meta = BlockApprovalMeta {
@@ -353,7 +412,7 @@ fn try_import_the_same_assignment() {
 		let assignments = vec![(cert.clone(), 0u32)];
 
 		let msg = protocol_v1::ApprovalDistributionMessage::Assignments(assignments.clone());
-		send_message_from_peer(overseer, &peer_a, Versioned::V1(msg)).await;
+		send_message_from_peer(overseer, &peer_a, msg).await;
 
 		expect_reputation_change(overseer, &peer_a, COST_UNEXPECTED_MESSAGE).await;
 
@@ -362,10 +421,11 @@ fn try_import_the_same_assignment() {
 			overseer_recv(overseer).await,
 			AllMessages::ApprovalVoting(ApprovalVotingMessage::CheckAndImportAssignment(
 				assignment,
-				0u32,
+				claimed_indices,
 				tx,
 			)) => {
-				assert_eq!(assignment, cert);
+				assert_eq!(claimed_indices, 0u32.into());
+				assert_eq!(assignment, cert.into());
 				tx.send(AssignmentCheckResult::Accepted).unwrap();
 			}
 		);
@@ -385,12 +445,104 @@ fn try_import_the_same_assignment() {
 			}
 		);
 
-		// setup new peer
-		setup_peer_with_view(overseer, &peer_d, ValidationVersion::V1, view![]).await;
+		// setup new peer with V2
+		setup_peer_with_view(overseer, &peer_d, view![], ValidationVersion::VStaging).await;
 
 		// send the same assignment from peer_d
 		let msg = protocol_v1::ApprovalDistributionMessage::Assignments(assignments);
-		send_message_from_peer(overseer, &peer_d, Versioned::V1(msg)).await;
+		send_message_from_peer(overseer, &peer_d, msg).await;
+
+		expect_reputation_change(overseer, &peer_d, COST_UNEXPECTED_MESSAGE).await;
+		expect_reputation_change(overseer, &peer_d, BENEFIT_VALID_MESSAGE).await;
+
+		assert!(overseer.recv().timeout(TIMEOUT).await.is_none(), "no message should be sent");
+		virtual_overseer
+	});
+}
+
+/// Just like `try_import_the_same_assignment` but use `VRFModuloCompact` assignments for multiple
+/// cores.
+#[test]
+fn try_import_the_same_assignment_v2() {
+	let peer_a = PeerId::random();
+	let peer_b = PeerId::random();
+	let peer_c = PeerId::random();
+	let peer_d = PeerId::random();
+	let parent_hash = Hash::repeat_byte(0xFF);
+	let hash = Hash::repeat_byte(0xAA);
+
+	let _ = test_harness(state_without_reputation_delay(), |mut virtual_overseer| async move {
+		let overseer = &mut virtual_overseer;
+		// setup peers
+		setup_peer_with_view(overseer, &peer_a, view![], ValidationVersion::VStaging).await;
+		setup_peer_with_view(overseer, &peer_b, view![hash], ValidationVersion::VStaging).await;
+		setup_peer_with_view(overseer, &peer_c, view![hash], ValidationVersion::VStaging).await;
+
+		// new block `hash_a` with 1 candidates
+		let meta = BlockApprovalMeta {
+			hash,
+			parent_hash,
+			number: 2,
+			candidates: vec![Default::default(); 1],
+			slot: 1.into(),
+			session: 1,
+		};
+		let msg = ApprovalDistributionMessage::NewBlocks(vec![meta]);
+		overseer_send(overseer, msg).await;
+
+		// send the assignment related to `hash`
+		let validator_index = ValidatorIndex(0);
+		let cores = vec![1, 2, 3, 4];
+		let core_bitfield: CoreBitfield = cores
+			.iter()
+			.map(|index| CoreIndex(*index))
+			.collect::<Vec<_>>()
+			.try_into()
+			.unwrap();
+
+		let cert = fake_assignment_cert_v2(hash, validator_index, core_bitfield.clone());
+		let assignments = vec![(cert.clone(), cores.clone().try_into().unwrap())];
+
+		let msg = protocol_vstaging::ApprovalDistributionMessage::Assignments(assignments.clone());
+		send_message_from_peer_vstaging(overseer, &peer_a, msg).await;
+
+		expect_reputation_change(overseer, &peer_a, COST_UNEXPECTED_MESSAGE).await;
+
+		// send an `Accept` message from the Approval Voting subsystem
+		assert_matches!(
+			overseer_recv(overseer).await,
+			AllMessages::ApprovalVoting(ApprovalVotingMessage::CheckAndImportAssignment(
+				assignment,
+				claimed_indices,
+				tx,
+			)) => {
+				assert_eq!(claimed_indices, cores.try_into().unwrap());
+				assert_eq!(assignment, cert.into());
+				tx.send(AssignmentCheckResult::Accepted).unwrap();
+			}
+		);
+
+		expect_reputation_change(overseer, &peer_a, BENEFIT_VALID_MESSAGE_FIRST).await;
+
+		assert_matches!(
+			overseer_recv(overseer).await,
+			AllMessages::NetworkBridgeTx(NetworkBridgeTxMessage::SendValidationMessage(
+				peers,
+				Versioned::VStaging(protocol_vstaging::ValidationProtocol::ApprovalDistribution(
+					protocol_vstaging::ApprovalDistributionMessage::Assignments(assignments)
+				))
+			)) => {
+				assert_eq!(peers.len(), 2);
+				assert_eq!(assignments.len(), 1);
+			}
+		);
+
+		// setup new peer
+		setup_peer_with_view(overseer, &peer_d, view![], ValidationVersion::VStaging).await;
+
+		// send the same assignment from peer_d
+		let msg = protocol_vstaging::ApprovalDistributionMessage::Assignments(assignments);
+		send_message_from_peer_vstaging(overseer, &peer_d, msg).await;
 
 		expect_reputation_change(overseer, &peer_d, COST_UNEXPECTED_MESSAGE).await;
 		expect_reputation_change(overseer, &peer_d, BENEFIT_VALID_MESSAGE).await;
@@ -413,7 +565,7 @@ fn delay_reputation_change() {
 		let overseer = &mut virtual_overseer;
 
 		// Setup peers
-		setup_peer_with_view(overseer, &peer, ValidationVersion::V1, view![]).await;
+		setup_peer_with_view(overseer, &peer, view![], ValidationVersion::V1).await;
 
 		// new block `hash_a` with 1 candidates
 		let meta = BlockApprovalMeta {
@@ -433,17 +585,18 @@ fn delay_reputation_change() {
 		let assignments = vec![(cert.clone(), 0u32)];
 
 		let msg = protocol_v1::ApprovalDistributionMessage::Assignments(assignments.clone());
-		send_message_from_peer(overseer, &peer, Versioned::V1(msg)).await;
+		send_message_from_peer(overseer, &peer, msg).await;
 
 		// send an `Accept` message from the Approval Voting subsystem
 		assert_matches!(
 			overseer_recv(overseer).await,
 			AllMessages::ApprovalVoting(ApprovalVotingMessage::CheckAndImportAssignment(
 				assignment,
-				0u32,
+				claimed_candidates,
 				tx,
 			)) => {
-				assert_eq!(assignment, cert);
+				assert_eq!(assignment.cert, cert.cert.into());
+				assert_eq!(claimed_candidates, vec![0u32].try_into().unwrap());
 				tx.send(AssignmentCheckResult::Accepted).unwrap();
 			}
 		);
@@ -474,7 +627,7 @@ fn spam_attack_results_in_negative_reputation_change() {
 	let _ = test_harness(state_without_reputation_delay(), |mut virtual_overseer| async move {
 		let overseer = &mut virtual_overseer;
 		let peer = &peer_a;
-		setup_peer_with_view(overseer, peer, ValidationVersion::V1, view![]).await;
+		setup_peer_with_view(overseer, peer, view![], ValidationVersion::V1).await;
 
 		// new block `hash_b` with 20 candidates
 		let candidates_count = 20;
@@ -501,7 +654,7 @@ fn spam_attack_results_in_negative_reputation_change() {
 			.collect();
 
 		let msg = protocol_v1::ApprovalDistributionMessage::Assignments(assignments.clone());
-		send_message_from_peer(overseer, peer, Versioned::V1(msg.clone())).await;
+		send_message_from_peer(overseer, peer, msg.clone()).await;
 
 		for i in 0..candidates_count {
 			expect_reputation_change(overseer, peer, COST_UNEXPECTED_MESSAGE).await;
@@ -513,8 +666,8 @@ fn spam_attack_results_in_negative_reputation_change() {
 					claimed_candidate_index,
 					tx,
 				)) => {
-					assert_eq!(assignment, assignments[i].0);
-					assert_eq!(claimed_candidate_index, assignments[i].1);
+					assert_eq!(assignment, assignments[i].0.clone().into());
+					assert_eq!(claimed_candidate_index, assignments[i].1.into());
 					tx.send(AssignmentCheckResult::Accepted).unwrap();
 				}
 			);
@@ -533,7 +686,7 @@ fn spam_attack_results_in_negative_reputation_change() {
 		.await;
 
 		// send the assignments again
-		send_message_from_peer(overseer, peer, Versioned::V1(msg.clone())).await;
+		send_message_from_peer(overseer, peer, msg.clone()).await;
 
 		// each of them will incur `COST_UNEXPECTED_MESSAGE`, not only the first one
 		for _ in 0..candidates_count {
@@ -558,7 +711,7 @@ fn peer_sending_us_the_same_we_just_sent_them_is_ok() {
 	let _ = test_harness(state_without_reputation_delay(), |mut virtual_overseer| async move {
 		let overseer = &mut virtual_overseer;
 		let peer = &peer_a;
-		setup_peer_with_view(overseer, peer, ValidationVersion::V1, view![]).await;
+		setup_peer_with_view(overseer, peer, view![], ValidationVersion::V1).await;
 
 		// new block `hash` with 1 candidates
 		let meta = BlockApprovalMeta {
@@ -578,7 +731,10 @@ fn peer_sending_us_the_same_we_just_sent_them_is_ok() {
 		let cert = fake_assignment_cert(hash, validator_index);
 		overseer_send(
 			overseer,
-			ApprovalDistributionMessage::DistributeAssignment(cert.clone(), candidate_index),
+			ApprovalDistributionMessage::DistributeAssignment(
+				cert.clone().into(),
+				candidate_index.into(),
+			),
 		)
 		.await;
 
@@ -610,12 +766,12 @@ fn peer_sending_us_the_same_we_just_sent_them_is_ok() {
 		// the peer could send us it as well
 		let assignments = vec![(cert, candidate_index)];
 		let msg = protocol_v1::ApprovalDistributionMessage::Assignments(assignments);
-		send_message_from_peer(overseer, peer, Versioned::V1(msg.clone())).await;
+		send_message_from_peer(overseer, peer, msg.clone()).await;
 
 		assert!(overseer.recv().timeout(TIMEOUT).await.is_none(), "we should not punish the peer");
 
 		// send the assignments again
-		send_message_from_peer(overseer, peer, Versioned::V1(msg)).await;
+		send_message_from_peer(overseer, peer, msg).await;
 
 		// now we should
 		expect_reputation_change(overseer, peer, COST_DUPLICATE_MESSAGE).await;
@@ -633,10 +789,10 @@ fn import_approval_happy_path() {
 
 	let _ = test_harness(state_without_reputation_delay(), |mut virtual_overseer| async move {
 		let overseer = &mut virtual_overseer;
-		// setup peers
-		setup_peer_with_view(overseer, &peer_a, ValidationVersion::V1, view![]).await;
-		setup_peer_with_view(overseer, &peer_b, ValidationVersion::V1, view![hash]).await;
-		setup_peer_with_view(overseer, &peer_c, ValidationVersion::V1, view![hash]).await;
+		// setup peers with V1 and V2 protocol versions
+		setup_peer_with_view(overseer, &peer_a, view![], ValidationVersion::V1).await;
+		setup_peer_with_view(overseer, &peer_b, view![hash], ValidationVersion::VStaging).await;
+		setup_peer_with_view(overseer, &peer_c, view![hash], ValidationVersion::V1).await;
 
 		// new block `hash_a` with 1 candidates
 		let meta = BlockApprovalMeta {
@@ -656,10 +812,14 @@ fn import_approval_happy_path() {
 		let cert = fake_assignment_cert(hash, validator_index);
 		overseer_send(
 			overseer,
-			ApprovalDistributionMessage::DistributeAssignment(cert, candidate_index),
+			ApprovalDistributionMessage::DistributeAssignment(
+				cert.clone().into(),
+				candidate_index.into(),
+			),
 		)
 		.await;
 
+		// 1 peer is v1
 		assert_matches!(
 			overseer_recv(overseer).await,
 			AllMessages::NetworkBridgeTx(NetworkBridgeTxMessage::SendValidationMessage(
@@ -668,7 +828,21 @@ fn import_approval_happy_path() {
 					protocol_v1::ApprovalDistributionMessage::Assignments(assignments)
 				))
 			)) => {
-				assert_eq!(peers.len(), 2);
+				assert_eq!(peers.len(), 1);
+				assert_eq!(assignments.len(), 1);
+			}
+		);
+
+		// 1 peer is v2
+		assert_matches!(
+			overseer_recv(overseer).await,
+			AllMessages::NetworkBridgeTx(NetworkBridgeTxMessage::SendValidationMessage(
+				peers,
+				Versioned::VStaging(protocol_vstaging::ValidationProtocol::ApprovalDistribution(
+					protocol_vstaging::ApprovalDistributionMessage::Assignments(assignments)
+				))
+			)) => {
+				assert_eq!(peers.len(), 1);
 				assert_eq!(assignments.len(), 1);
 			}
 		);
@@ -681,7 +855,7 @@ fn import_approval_happy_path() {
 			signature: dummy_signature(),
 		};
 		let msg = protocol_v1::ApprovalDistributionMessage::Approvals(vec![approval.clone()]);
-		send_message_from_peer(overseer, &peer_b, Versioned::V1(msg)).await;
+		send_message_from_peer(overseer, &peer_b, msg).await;
 
 		assert_matches!(
 			overseer_recv(overseer).await,
@@ -722,8 +896,8 @@ fn import_approval_bad() {
 	let _ = test_harness(state_without_reputation_delay(), |mut virtual_overseer| async move {
 		let overseer = &mut virtual_overseer;
 		// setup peers
-		setup_peer_with_view(overseer, &peer_a, ValidationVersion::V1, view![]).await;
-		setup_peer_with_view(overseer, &peer_b, ValidationVersion::V1, view![hash]).await;
+		setup_peer_with_view(overseer, &peer_a, view![], ValidationVersion::V1).await;
+		setup_peer_with_view(overseer, &peer_b, view![hash], ValidationVersion::V1).await;
 
 		// new block `hash_a` with 1 candidates
 		let meta = BlockApprovalMeta {
@@ -749,14 +923,14 @@ fn import_approval_bad() {
 			signature: dummy_signature(),
 		};
 		let msg = protocol_v1::ApprovalDistributionMessage::Approvals(vec![approval.clone()]);
-		send_message_from_peer(overseer, &peer_b, Versioned::V1(msg)).await;
+		send_message_from_peer(overseer, &peer_b, msg).await;
 
 		expect_reputation_change(overseer, &peer_b, COST_UNEXPECTED_MESSAGE).await;
 
 		// now import an assignment from peer_b
 		let assignments = vec![(cert.clone(), candidate_index)];
 		let msg = protocol_v1::ApprovalDistributionMessage::Assignments(assignments);
-		send_message_from_peer(overseer, &peer_b, Versioned::V1(msg)).await;
+		send_message_from_peer(overseer, &peer_b, msg).await;
 
 		assert_matches!(
 			overseer_recv(overseer).await,
@@ -765,8 +939,8 @@ fn import_approval_bad() {
 				i,
 				tx,
 			)) => {
-				assert_eq!(assignment, cert);
-				assert_eq!(i, candidate_index);
+				assert_eq!(assignment, cert.into());
+				assert_eq!(i, candidate_index.into());
 				tx.send(AssignmentCheckResult::Accepted).unwrap();
 			}
 		);
@@ -775,7 +949,7 @@ fn import_approval_bad() {
 
 		// and try again
 		let msg = protocol_v1::ApprovalDistributionMessage::Approvals(vec![approval.clone()]);
-		send_message_from_peer(overseer, &peer_b, Versioned::V1(msg)).await;
+		send_message_from_peer(overseer, &peer_b, msg).await;
 
 		assert_matches!(
 			overseer_recv(overseer).await,
@@ -911,12 +1085,20 @@ fn update_peer_view() {
 		let cert_a = fake_assignment_cert(hash_a, ValidatorIndex(0));
 		let cert_b = fake_assignment_cert(hash_b, ValidatorIndex(0));
 
-		overseer_send(overseer, ApprovalDistributionMessage::DistributeAssignment(cert_a, 0)).await;
+		overseer_send(
+			overseer,
+			ApprovalDistributionMessage::DistributeAssignment(cert_a.into(), 0.into()),
+		)
+		.await;
 
-		overseer_send(overseer, ApprovalDistributionMessage::DistributeAssignment(cert_b, 0)).await;
+		overseer_send(
+			overseer,
+			ApprovalDistributionMessage::DistributeAssignment(cert_b.into(), 0.into()),
+		)
+		.await;
 
 		// connect a peer
-		setup_peer_with_view(overseer, peer, ValidationVersion::V1, view![hash_a]).await;
+		setup_peer_with_view(overseer, peer, view![hash_a], ValidationVersion::V1).await;
 
 		// we should send relevant assignments to the peer
 		assert_matches!(
@@ -934,7 +1116,7 @@ fn update_peer_view() {
 		virtual_overseer
 	});
 
-	assert_eq!(state.peer_data.get(peer).map(|data| data.view.finalized_number), Some(0));
+	assert_eq!(state.peer_views.get(peer).map(|v| v.view.finalized_number), Some(0));
 	assert_eq!(
 		state
 			.blocks
@@ -965,7 +1147,7 @@ fn update_peer_view() {
 
 		overseer_send(
 			overseer,
-			ApprovalDistributionMessage::DistributeAssignment(cert_c.clone(), 0),
+			ApprovalDistributionMessage::DistributeAssignment(cert_c.clone().into(), 0.into()),
 		)
 		.await;
 
@@ -986,7 +1168,7 @@ fn update_peer_view() {
 		virtual_overseer
 	});
 
-	assert_eq!(state.peer_data.get(peer).map(|data| data.view.finalized_number), Some(2));
+	assert_eq!(state.peer_views.get(peer).map(|v| v.view.finalized_number), Some(2));
 	assert_eq!(
 		state
 			.blocks
@@ -1016,10 +1198,7 @@ fn update_peer_view() {
 		virtual_overseer
 	});
 
-	assert_eq!(
-		state.peer_data.get(peer).map(|data| data.view.finalized_number),
-		Some(finalized_number)
-	);
+	assert_eq!(state.peer_views.get(peer).map(|v| v.view.finalized_number), Some(finalized_number));
 	assert!(state.blocks.get(&hash_c).unwrap().known_by.get(peer).is_none());
 }
 
@@ -1034,7 +1213,7 @@ fn import_remotely_then_locally() {
 	let _ = test_harness(state_without_reputation_delay(), |mut virtual_overseer| async move {
 		let overseer = &mut virtual_overseer;
 		// setup the peer
-		setup_peer_with_view(overseer, peer, ValidationVersion::V1, view![hash]).await;
+		setup_peer_with_view(overseer, peer, view![hash], ValidationVersion::V1).await;
 
 		// new block `hash_a` with 1 candidates
 		let meta = BlockApprovalMeta {
@@ -1054,7 +1233,7 @@ fn import_remotely_then_locally() {
 		let cert = fake_assignment_cert(hash, validator_index);
 		let assignments = vec![(cert.clone(), candidate_index)];
 		let msg = protocol_v1::ApprovalDistributionMessage::Assignments(assignments.clone());
-		send_message_from_peer(overseer, peer, Versioned::V1(msg)).await;
+		send_message_from_peer(overseer, peer, msg).await;
 
 		// send an `Accept` message from the Approval Voting subsystem
 		assert_matches!(
@@ -1064,8 +1243,8 @@ fn import_remotely_then_locally() {
 				i,
 				tx,
 			)) => {
-				assert_eq!(assignment, cert);
-				assert_eq!(i, candidate_index);
+				assert_eq!(assignment, cert.clone().into());
+				assert_eq!(i, candidate_index.into());
 				tx.send(AssignmentCheckResult::Accepted).unwrap();
 			}
 		);
@@ -1075,7 +1254,10 @@ fn import_remotely_then_locally() {
 		// import the same assignment locally
 		overseer_send(
 			overseer,
-			ApprovalDistributionMessage::DistributeAssignment(cert, candidate_index),
+			ApprovalDistributionMessage::DistributeAssignment(
+				cert.clone().into(),
+				candidate_index.into(),
+			),
 		)
 		.await;
 
@@ -1089,7 +1271,7 @@ fn import_remotely_then_locally() {
 			signature: dummy_signature(),
 		};
 		let msg = protocol_v1::ApprovalDistributionMessage::Approvals(vec![approval.clone()]);
-		send_message_from_peer(overseer, peer, Versioned::V1(msg)).await;
+		send_message_from_peer(overseer, peer, msg).await;
 
 		assert_matches!(
 			overseer_recv(overseer).await,
@@ -1147,7 +1329,10 @@ fn sends_assignments_even_when_state_is_approved() {
 
 		overseer_send(
 			overseer,
-			ApprovalDistributionMessage::DistributeAssignment(cert.clone(), candidate_index),
+			ApprovalDistributionMessage::DistributeAssignment(
+				cert.clone().into(),
+				candidate_index.into(),
+			),
 		)
 		.await;
 
@@ -1155,7 +1340,7 @@ fn sends_assignments_even_when_state_is_approved() {
 			.await;
 
 		// connect the peer.
-		setup_peer_with_view(overseer, peer, ValidationVersion::V1, view![hash]).await;
+		setup_peer_with_view(overseer, peer, view![hash], ValidationVersion::V1).await;
 
 		let assignments = vec![(cert.clone(), candidate_index)];
 		let approvals = vec![approval.clone()];
@@ -1191,6 +1376,112 @@ fn sends_assignments_even_when_state_is_approved() {
 	});
 }
 
+/// Same as `sends_assignments_even_when_state_is_approved_v2` but with `VRFModuloCompact`
+/// assignemnts.
+#[test]
+fn sends_assignments_even_when_state_is_approved_v2() {
+	let peer_a = PeerId::random();
+	let parent_hash = Hash::repeat_byte(0xFF);
+	let hash = Hash::repeat_byte(0xAA);
+	let peer = &peer_a;
+
+	let _ = test_harness(State::default(), |mut virtual_overseer| async move {
+		let overseer = &mut virtual_overseer;
+
+		// new block `hash_a` with 1 candidates
+		let meta = BlockApprovalMeta {
+			hash,
+			parent_hash,
+			number: 1,
+			candidates: vec![Default::default(); 4],
+			slot: 1.into(),
+			session: 1,
+		};
+		let msg = ApprovalDistributionMessage::NewBlocks(vec![meta]);
+		overseer_send(overseer, msg).await;
+
+		let validator_index = ValidatorIndex(0);
+		let cores = vec![0, 1, 2, 3];
+		let candidate_bitfield: CandidateBitfield = cores.clone().try_into().unwrap();
+
+		let core_bitfield: CoreBitfield = cores
+			.iter()
+			.map(|index| CoreIndex(*index))
+			.collect::<Vec<_>>()
+			.try_into()
+			.unwrap();
+
+		let cert = fake_assignment_cert_v2(hash, validator_index, core_bitfield.clone());
+
+		// Assumes candidate index == core index.
+		let approvals = cores
+			.iter()
+			.map(|core| IndirectSignedApprovalVote {
+				block_hash: hash,
+				candidate_index: *core,
+				validator: validator_index,
+				signature: dummy_signature(),
+			})
+			.collect::<Vec<_>>();
+
+		overseer_send(
+			overseer,
+			ApprovalDistributionMessage::DistributeAssignment(
+				cert.clone().into(),
+				candidate_bitfield.clone(),
+			),
+		)
+		.await;
+
+		for approval in &approvals {
+			overseer_send(
+				overseer,
+				ApprovalDistributionMessage::DistributeApproval(approval.clone()),
+			)
+			.await;
+		}
+
+		// connect the peer.
+		setup_peer_with_view(overseer, peer, view![hash], ValidationVersion::VStaging).await;
+
+		let assignments = vec![(cert.clone(), candidate_bitfield.clone())];
+
+		assert_matches!(
+			overseer_recv(overseer).await,
+			AllMessages::NetworkBridgeTx(NetworkBridgeTxMessage::SendValidationMessage(
+				peers,
+				Versioned::VStaging(protocol_vstaging::ValidationProtocol::ApprovalDistribution(
+					protocol_vstaging::ApprovalDistributionMessage::Assignments(sent_assignments)
+				))
+			)) => {
+				assert_eq!(peers, vec![*peer]);
+				assert_eq!(sent_assignments, assignments);
+			}
+		);
+
+		assert_matches!(
+			overseer_recv(overseer).await,
+			AllMessages::NetworkBridgeTx(NetworkBridgeTxMessage::SendValidationMessage(
+				peers,
+				Versioned::VStaging(protocol_vstaging::ValidationProtocol::ApprovalDistribution(
+					protocol_vstaging::ApprovalDistributionMessage::Approvals(sent_approvals)
+				))
+			)) => {
+				// Construct a hashmaps of approvals for comparison. Approval distribution reorders messages because they are kept in a
+				// hashmap as well.
+				let sent_approvals = sent_approvals.into_iter().map(|approval| (approval.candidate_index, approval)).collect::<HashMap<_,_>>();
+				let approvals = approvals.into_iter().map(|approval| (approval.candidate_index, approval)).collect::<HashMap<_,_>>();
+
+				assert_eq!(peers, vec![*peer]);
+				assert_eq!(sent_approvals, approvals);
+			}
+		);
+
+		assert!(overseer.recv().timeout(TIMEOUT).await.is_none(), "no message should be sent");
+		virtual_overseer
+	});
+}
+
 /// <https://github.com/paritytech/polkadot/pull/5089>
 ///
 /// 1. Receive remote peer view update with an unknown head
@@ -1219,7 +1510,7 @@ fn race_condition_in_local_vs_remote_view_update() {
 		};
 
 		// This will send a peer view that is ahead of our view
-		setup_peer_with_view(overseer, peer, ValidationVersion::V1, view![hash_b]).await;
+		setup_peer_with_view(overseer, peer, view![hash_b], ValidationVersion::V1).await;
 
 		// Send our view update to include a new head
 		overseer_send(
@@ -1240,7 +1531,7 @@ fn race_condition_in_local_vs_remote_view_update() {
 			.collect();
 
 		let msg = protocol_v1::ApprovalDistributionMessage::Assignments(assignments.clone());
-		send_message_from_peer(overseer, peer, Versioned::V1(msg.clone())).await;
+		send_message_from_peer(overseer, peer, msg.clone()).await;
 
 		// This will handle pending messages being processed
 		let msg = ApprovalDistributionMessage::NewBlocks(vec![meta]);
@@ -1257,8 +1548,8 @@ fn race_condition_in_local_vs_remote_view_update() {
 					claimed_candidate_index,
 					tx,
 				)) => {
-					assert_eq!(assignment, assignments[i].0);
-					assert_eq!(claimed_candidate_index, assignments[i].1);
+					assert_eq!(assignment, assignments[i].0.clone().into());
+					assert_eq!(claimed_candidate_index, assignments[i].1.into());
 					tx.send(AssignmentCheckResult::Accepted).unwrap();
 				}
 			);
@@ -1283,7 +1574,7 @@ fn propagates_locally_generated_assignment_to_both_dimensions() {
 
 		// Connect all peers.
 		for (peer, _) in &peers {
-			setup_peer_with_view(overseer, peer, ValidationVersion::V1, view![hash]).await;
+			setup_peer_with_view(overseer, peer, view![hash], ValidationVersion::V1).await;
 		}
 
 		// Set up a gossip topology.
@@ -1325,7 +1616,10 @@ fn propagates_locally_generated_assignment_to_both_dimensions() {
 
 		overseer_send(
 			overseer,
-			ApprovalDistributionMessage::DistributeAssignment(cert.clone(), candidate_index),
+			ApprovalDistributionMessage::DistributeAssignment(
+				cert.clone().into(),
+				candidate_index.into(),
+			),
 		)
 		.await;
 
@@ -1388,7 +1682,7 @@ fn propagates_assignments_along_unshared_dimension() {
 
 		// Connect all peers.
 		for (peer, _) in &peers {
-			setup_peer_with_view(overseer, peer, ValidationVersion::V1, view![hash]).await;
+			setup_peer_with_view(overseer, peer, view![hash], ValidationVersion::V1).await;
 		}
 
 		// Set up a gossip topology.
@@ -1424,7 +1718,7 @@ fn propagates_assignments_along_unshared_dimension() {
 
 			// Issuer of the message is important, not the peer we receive from.
 			// 99 deliberately chosen because it's not in X or Y.
-			send_message_from_peer(overseer, &peers[99].0, Versioned::V1(msg)).await;
+			send_message_from_peer(overseer, &peers[99].0, msg).await;
 			assert_matches!(
 				overseer_recv(overseer).await,
 				AllMessages::ApprovalVoting(ApprovalVotingMessage::CheckAndImportAssignment(
@@ -1473,7 +1767,7 @@ fn propagates_assignments_along_unshared_dimension() {
 
 			// Issuer of the message is important, not the peer we receive from.
 			// 99 deliberately chosen because it's not in X or Y.
-			send_message_from_peer(overseer, &peers[99].0, Versioned::V1(msg)).await;
+			send_message_from_peer(overseer, &peers[99].0, msg).await;
 			assert_matches!(
 				overseer_recv(overseer).await,
 				AllMessages::ApprovalVoting(ApprovalVotingMessage::CheckAndImportAssignment(
@@ -1530,7 +1824,7 @@ fn propagates_to_required_after_connect() {
 		// Connect all peers except omitted.
 		for (i, (peer, _)) in peers.iter().enumerate() {
 			if !omitted.contains(&i) {
-				setup_peer_with_view(overseer, peer, ValidationVersion::V1, view![hash]).await;
+				setup_peer_with_view(overseer, peer, view![hash], ValidationVersion::V1).await;
 			}
 		}
 
@@ -1573,7 +1867,10 @@ fn propagates_to_required_after_connect() {
 
 		overseer_send(
 			overseer,
-			ApprovalDistributionMessage::DistributeAssignment(cert.clone(), candidate_index),
+			ApprovalDistributionMessage::DistributeAssignment(
+				cert.clone().into(),
+				candidate_index.into(),
+			),
 		)
 		.await;
 
@@ -1619,7 +1916,7 @@ fn propagates_to_required_after_connect() {
 		);
 
 		for i in omitted.iter().copied() {
-			setup_peer_with_view(overseer, &peers[i].0, ValidationVersion::V1, view![hash]).await;
+			setup_peer_with_view(overseer, &peers[i].0, view![hash], ValidationVersion::V1).await;
 
 			assert_matches!(
 				overseer_recv(overseer).await,
@@ -1668,7 +1965,7 @@ fn sends_to_more_peers_after_getting_topology() {
 
 		// Connect all peers except omitted.
 		for (peer, _) in &peers {
-			setup_peer_with_view(overseer, peer, ValidationVersion::V1, view![hash]).await;
+			setup_peer_with_view(overseer, peer, view![hash], ValidationVersion::V1).await;
 		}
 
 		// new block `hash_a` with 1 candidates
@@ -1698,7 +1995,10 @@ fn sends_to_more_peers_after_getting_topology() {
 
 		overseer_send(
 			overseer,
-			ApprovalDistributionMessage::DistributeAssignment(cert.clone(), candidate_index),
+			ApprovalDistributionMessage::DistributeAssignment(
+				cert.clone().into(),
+				candidate_index.into(),
+			),
 		)
 		.await;
 
@@ -1820,7 +2120,7 @@ fn originator_aggression_l1() {
 
 		// Connect all peers except omitted.
 		for (peer, _) in &peers {
-			setup_peer_with_view(overseer, peer, ValidationVersion::V1, view![hash]).await;
+			setup_peer_with_view(overseer, peer, view![hash], ValidationVersion::V1).await;
 		}
 
 		// new block `hash_a` with 1 candidates
@@ -1857,7 +2157,10 @@ fn originator_aggression_l1() {
 
 		overseer_send(
 			overseer,
-			ApprovalDistributionMessage::DistributeAssignment(cert.clone(), candidate_index),
+			ApprovalDistributionMessage::DistributeAssignment(
+				cert.clone().into(),
+				candidate_index.into(),
+			),
 		)
 		.await;
 
@@ -1979,7 +2282,7 @@ fn non_originator_aggression_l1() {
 
 		// Connect all peers except omitted.
 		for (peer, _) in &peers {
-			setup_peer_with_view(overseer, peer, ValidationVersion::V1, view![hash]).await;
+			setup_peer_with_view(overseer, peer, view![hash], ValidationVersion::V1).await;
 		}
 
 		// new block `hash_a` with 1 candidates
@@ -2008,12 +2311,12 @@ fn non_originator_aggression_l1() {
 		)
 		.await;
 
-		let assignments = vec![(cert.clone(), candidate_index)];
+		let assignments = vec![(cert.clone().into(), candidate_index)];
 		let msg = protocol_v1::ApprovalDistributionMessage::Assignments(assignments.clone());
 
 		// Issuer of the message is important, not the peer we receive from.
 		// 99 deliberately chosen because it's not in X or Y.
-		send_message_from_peer(overseer, &peers[99].0, Versioned::V1(msg)).await;
+		send_message_from_peer(overseer, &peers[99].0, msg).await;
 		assert_matches!(
 			overseer_recv(overseer).await,
 			AllMessages::ApprovalVoting(ApprovalVotingMessage::CheckAndImportAssignment(
@@ -2084,7 +2387,7 @@ fn non_originator_aggression_l2() {
 
 		// Connect all peers except omitted.
 		for (peer, _) in &peers {
-			setup_peer_with_view(overseer, peer, ValidationVersion::V1, view![hash]).await;
+			setup_peer_with_view(overseer, peer, view![hash], ValidationVersion::V1).await;
 		}
 
 		// new block `hash_a` with 1 candidates
@@ -2118,7 +2421,7 @@ fn non_originator_aggression_l2() {
 
 		// Issuer of the message is important, not the peer we receive from.
 		// 99 deliberately chosen because it's not in X or Y.
-		send_message_from_peer(overseer, &peers[99].0, Versioned::V1(msg)).await;
+		send_message_from_peer(overseer, &peers[99].0, msg).await;
 		assert_matches!(
 			overseer_recv(overseer).await,
 			AllMessages::ApprovalVoting(ApprovalVotingMessage::CheckAndImportAssignment(
@@ -2249,7 +2552,7 @@ fn resends_messages_periodically() {
 
 		// Connect all peers.
 		for (peer, _) in &peers {
-			setup_peer_with_view(overseer, peer, ValidationVersion::V1, view![hash]).await;
+			setup_peer_with_view(overseer, peer, view![hash], ValidationVersion::V1).await;
 		}
 
 		// Set up a gossip topology.
@@ -2284,7 +2587,7 @@ fn resends_messages_periodically() {
 
 			// Issuer of the message is important, not the peer we receive from.
 			// 99 deliberately chosen because it's not in X or Y.
-			send_message_from_peer(overseer, &peers[99].0, Versioned::V1(msg)).await;
+			send_message_from_peer(overseer, &peers[99].0, msg).await;
 			assert_matches!(
 				overseer_recv(overseer).await,
 				AllMessages::ApprovalVoting(ApprovalVotingMessage::CheckAndImportAssignment(
@@ -2388,9 +2691,9 @@ fn import_versioned_approval() {
 	let _ = test_harness(state, |mut virtual_overseer| async move {
 		let overseer = &mut virtual_overseer;
 		// All peers are aware of relay parent.
-		setup_peer_with_view(overseer, &peer_a, ValidationVersion::V2, view![hash]).await;
-		setup_peer_with_view(overseer, &peer_b, ValidationVersion::V1, view![hash]).await;
-		setup_peer_with_view(overseer, &peer_c, ValidationVersion::V2, view![hash]).await;
+		setup_peer_with_view(overseer, &peer_a, view![hash], ValidationVersion::V2).await;
+		setup_peer_with_view(overseer, &peer_b, view![hash], ValidationVersion::V1).await;
+		setup_peer_with_view(overseer, &peer_c, view![hash], ValidationVersion::V2).await;
 
 		// new block `hash_a` with 1 candidates
 		let meta = BlockApprovalMeta {
@@ -2410,7 +2713,7 @@ fn import_versioned_approval() {
 		let cert = fake_assignment_cert(hash, validator_index);
 		overseer_send(
 			overseer,
-			ApprovalDistributionMessage::DistributeAssignment(cert, candidate_index),
+			ApprovalDistributionMessage::DistributeAssignment(cert.into(), candidate_index.into()),
 		)
 		.await;
 
@@ -2451,7 +2754,7 @@ fn import_versioned_approval() {
 			signature: dummy_signature(),
 		};
 		let msg = protocol_v2::ApprovalDistributionMessage::Approvals(vec![approval.clone()]);
-		send_message_from_peer(overseer, &peer_a, Versioned::V2(msg)).await;
+		send_message_from_peer_v2(overseer, &peer_a, msg).await;
 
 		assert_matches!(
 			overseer_recv(overseer).await,
@@ -2512,7 +2815,9 @@ fn batch_test_round(message_count: usize) {
 		let validators = 0..message_count;
 		let assignments: Vec<_> = validators
 			.clone()
-			.map(|index| (fake_assignment_cert(Hash::zero(), ValidatorIndex(index as u32)), 0))
+			.map(|index| {
+				(fake_assignment_cert(Hash::zero(), ValidatorIndex(index as u32)).into(), 0.into())
+			})
 			.collect();
 
 		let approvals: Vec<_> = validators
@@ -2525,9 +2830,18 @@ fn batch_test_round(message_count: usize) {
 			.collect();
 
 		let peer = PeerId::random();
-		send_assignments_batched(&mut sender, assignments.clone(), peer, ValidationVersion::V1)
-			.await;
-		send_approvals_batched(&mut sender, approvals.clone(), peer, ValidationVersion::V1).await;
+		send_assignments_batched(
+			&mut sender,
+			assignments.clone(),
+			&vec![(peer, ValidationVersion::V1.into())],
+		)
+		.await;
+		send_approvals_batched(
+			&mut sender,
+			approvals.clone(),
+			&vec![(peer, ValidationVersion::V1.into())],
+		)
+		.await;
 
 		// Check expected assignments batches.
 		for assignment_index in (0..assignments.len()).step_by(super::MAX_ASSIGNMENT_BATCH_SIZE) {
@@ -2549,7 +2863,7 @@ fn batch_test_round(message_count: usize) {
 					assert_eq!(peers.len(), 1);
 
 					for (message_index,  assignment) in sent_assignments.iter().enumerate() {
-						assert_eq!(assignment.0, assignments[assignment_index + message_index].0);
+						assert_eq!(assignment.0, assignments[assignment_index + message_index].0.clone().try_into().unwrap());
 						assert_eq!(assignment.1, 0);
 					}
 				}