collation-genereation: fix tests (#3883)

Somehow https://github.com/paritytech/polkadot-sdk/pull/3795 was merged
but tests are failing now on master. I suspect that CI is not even
running these tests anymore which is a big issue.

---------

Signed-off-by: Andrei Sandu <andrei-mihail@parity.io>

polkadot/node/collation-generation/src/tests.rs

@@ -686,7 +686,7 @@ fn submit_collation_is_no_op_before_initialization() {
 fn submit_collation_leads_to_distribution() {
 	let relay_parent = Hash::repeat_byte(0);
 	let validation_code_hash = ValidationCodeHash::from(Hash::repeat_byte(42));
-	let parent_head = HeadData::from(vec![1, 2, 3]);
+	let parent_head = dummy_head_data();
 	let para_id = ParaId::from(5);
 	let expected_pvd = PersistedValidationData {
 		parent_head: parent_head.clone(),
@@ -707,7 +707,7 @@ fn submit_collation_leads_to_distribution() {
 				msg: CollationGenerationMessage::SubmitCollation(SubmitCollationParams {
 					relay_parent,
 					collation: test_collation(),
-					parent_head: vec![1, 2, 3].into(),
+					parent_head: dummy_head_data(),
 					validation_code_hash,
 					result_sender: None,
 					core_index: CoreIndex(0),
@@ -795,16 +795,16 @@ fn distribute_collation_for_occupied_core_with_async_backing_enabled(#[case] run
 			cores,
 			runtime_version,
 			claim_queue,
-			pending_availability,
 		)
 		.await;
-		helpers::handle_core_processing_for_a_leaf(
+		helpers::handle_cores_processing_for_a_leaf(
 			&mut virtual_overseer,
 			activated_hash,
 			para_id,
 			// `CoreState` is `Occupied` => `OccupiedCoreAssumption` is `Included`
 			OccupiedCoreAssumption::Included,
 			1,
+			pending_availability,
 		)
 		.await;
 
@@ -825,7 +825,7 @@ fn distribute_collation_for_occupied_cores_with_async_backing_enabled_and_elasti
 	let activated_hash: Hash = [1; 32].into();
 	let para_id = ParaId::from(5);
 
-	let cores = (0..candidates_pending_avail)
+	let cores = (0..3)
 		.into_iter()
 		.map(|idx| {
 			CoreState::Occupied(polkadot_primitives::OccupiedCore {
@@ -864,17 +864,22 @@ fn distribute_collation_for_occupied_cores_with_async_backing_enabled_and_elasti
 			// Using latest runtime with the fancy claim queue exposed.
 			RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT,
 			claim_queue,
-			pending_availability,
 		)
 		.await;
 
-		helpers::handle_core_processing_for_a_leaf(
+		helpers::handle_cores_processing_for_a_leaf(
 			&mut virtual_overseer,
 			activated_hash,
 			para_id,
-			// `CoreState` is `Occupied` => `OccupiedCoreAssumption` is `Included`
-			OccupiedCoreAssumption::Included,
+			// if at least 1 cores is occupied => `OccupiedCoreAssumption` is `Included`
+			// else assumption is `Free`.
+			if candidates_pending_avail > 0 {
+				OccupiedCoreAssumption::Included
+			} else {
+				OccupiedCoreAssumption::Free
+			},
 			total_cores,
+			pending_availability,
 		)
 		.await;
 
@@ -890,12 +895,12 @@ fn distribute_collation_for_occupied_cores_with_async_backing_enabled_and_elasti
 #[case(1)]
 #[case(2)]
 fn distribute_collation_for_free_cores_with_async_backing_enabled_and_elastic_scaling(
-	#[case] candidates_pending_avail: u32,
+	#[case] total_cores: usize,
 ) {
 	let activated_hash: Hash = [1; 32].into();
 	let para_id = ParaId::from(5);
 
-	let cores = (0..candidates_pending_avail)
+	let cores = (0..total_cores)
 		.into_iter()
 		.map(|_idx| CoreState::Scheduled(ScheduledCore { para_id, collator: None }))
 		.collect::<Vec<_>>();
@@ -905,7 +910,6 @@ fn distribute_collation_for_free_cores_with_async_backing_enabled_and_elastic_sc
 		.enumerate()
 		.map(|(idx, _core)| (CoreIndex::from(idx as u32), VecDeque::from([para_id])))
 		.collect::<BTreeMap<_, _>>();
-	let total_cores = cores.len();
 
 	test_harness(|mut virtual_overseer| async move {
 		helpers::initialize_collator(&mut virtual_overseer, para_id).await;
@@ -918,17 +922,17 @@ fn distribute_collation_for_free_cores_with_async_backing_enabled_and_elastic_sc
 			// Using latest runtime with the fancy claim queue exposed.
 			RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT,
 			claim_queue,
-			vec![],
 		)
 		.await;
 
-		helpers::handle_core_processing_for_a_leaf(
+		helpers::handle_cores_processing_for_a_leaf(
 			&mut virtual_overseer,
 			activated_hash,
 			para_id,
 			// `CoreState` is `Free` => `OccupiedCoreAssumption` is `Free`
 			OccupiedCoreAssumption::Free,
 			total_cores,
+			vec![],
 		)
 		.await;
 
@@ -963,6 +967,7 @@ fn no_collation_is_distributed_for_occupied_core_with_async_backing_disabled(
 	test_harness(|mut virtual_overseer| async move {
 		helpers::initialize_collator(&mut virtual_overseer, para_id).await;
 		helpers::activate_new_head(&mut virtual_overseer, activated_hash).await;
+
 		helpers::handle_runtime_calls_on_new_head_activation(
 			&mut virtual_overseer,
 			activated_hash,
@@ -970,7 +975,6 @@ fn no_collation_is_distributed_for_occupied_core_with_async_backing_disabled(
 			cores,
 			runtime_version,
 			claim_queue,
-			vec![],
 		)
 		.await;
 
@@ -1047,7 +1051,6 @@ mod helpers {
 		cores: Vec<CoreState>,
 		runtime_version: u32,
 		claim_queue: BTreeMap<CoreIndex, VecDeque<ParaId>>,
-		pending_availability: Vec<CandidatePendingAvailability>,
 	) {
 		assert_matches!(
 			overseer_recv(virtual_overseer).await,
@@ -1082,25 +1085,6 @@ mod helpers {
 			}
 		);
 
-		// Process the `ParaBackingState` message, and return some dummy state.
-		let message = overseer_recv(virtual_overseer).await;
-		let para_id = match message {
-			AllMessages::RuntimeApi(RuntimeApiMessage::Request(
-				_,
-				RuntimeApiRequest::ParaBackingState(p_id, _),
-			)) => p_id,
-			_ => panic!("received unexpected message {:?}", message),
-		};
-
-		assert_matches!(
-			message,
-			AllMessages::RuntimeApi(
-				RuntimeApiMessage::Request(parent, RuntimeApiRequest::ParaBackingState(p_id, tx))
-			) if parent == activated_hash && p_id == para_id => {
-				tx.send(Ok(Some(dummy_backing_state(pending_availability)))).unwrap();
-			}
-		);
-
 		assert_matches!(
 			overseer_recv(virtual_overseer).await,
 			AllMessages::RuntimeApi(RuntimeApiMessage::Request(
@@ -1128,12 +1112,13 @@ mod helpers {
 
 	// Handles all runtime requests performed in `handle_new_activations` for the case when a
 	// collation should be prepared for the new leaf
-	pub async fn handle_core_processing_for_a_leaf(
+	pub async fn handle_cores_processing_for_a_leaf(
 		virtual_overseer: &mut VirtualOverseer,
 		activated_hash: Hash,
 		para_id: ParaId,
 		expected_occupied_core_assumption: OccupiedCoreAssumption,
 		cores_assigned: usize,
+		pending_availability: Vec<CandidatePendingAvailability>,
 	) {
 		// Expect no messages if no cores is assigned to the para
 		if cores_assigned == 0 {
@@ -1143,7 +1128,7 @@ mod helpers {
 
 		// Some hardcoded data - if needed, extract to parameters
 		let validation_code_hash = ValidationCodeHash::from(Hash::repeat_byte(42));
-		let parent_head = HeadData::from(vec![1, 2, 3]);
+		let parent_head = dummy_head_data();
 		let pvd = PersistedValidationData {
 			parent_head: parent_head.clone(),
 			relay_parent_number: 10,
@@ -1151,6 +1136,15 @@ mod helpers {
 			max_pov_size: 1024,
 		};
 
+		assert_matches!(
+			overseer_recv(virtual_overseer).await,
+			AllMessages::RuntimeApi(
+				RuntimeApiMessage::Request(parent, RuntimeApiRequest::ParaBackingState(p_id, tx))
+			) if parent == activated_hash && p_id == para_id => {
+				tx.send(Ok(Some(dummy_backing_state(pending_availability)))).unwrap();
+			}
+		);
+
 		assert_matches!(
 			overseer_recv(virtual_overseer).await,
 			AllMessages::RuntimeApi(RuntimeApiMessage::Request(hash, RuntimeApiRequest::PersistedValidationData(id, a, tx))) => {
@@ -1180,18 +1174,20 @@ mod helpers {
 			}
 		);
 
-		assert_matches!(
-			overseer_recv(virtual_overseer).await,
-			AllMessages::CollatorProtocol(CollatorProtocolMessage::DistributeCollation{
-				candidate_receipt,
-				parent_head_data_hash,
-				..
-			}) => {
-				assert_eq!(parent_head_data_hash, parent_head.hash());
-				assert_eq!(candidate_receipt.descriptor().persisted_validation_data_hash, pvd.hash());
-				assert_eq!(candidate_receipt.descriptor().para_head, dummy_head_data().hash());
-				assert_eq!(candidate_receipt.descriptor().validation_code_hash, validation_code_hash);
-			}
-		);
+		for _ in 0..cores_assigned {
+			assert_matches!(
+				overseer_recv(virtual_overseer).await,
+				AllMessages::CollatorProtocol(CollatorProtocolMessage::DistributeCollation{
+					candidate_receipt,
+					parent_head_data_hash,
+					..
+				}) => {
+					assert_eq!(parent_head_data_hash, parent_head.hash());
+					assert_eq!(candidate_receipt.descriptor().persisted_validation_data_hash, pvd.hash());
+					assert_eq!(candidate_receipt.descriptor().para_head, dummy_head_data().hash());
+					assert_eq!(candidate_receipt.descriptor().validation_code_hash, validation_code_hash);
+				}
+			);
+		}
 	}
 }