Fix tests.

Author: eskimor

polkadot/node/network/collator-protocol/src/collator_side/tests/prospective_parachains.rs

@@ -75,21 +75,7 @@ pub(super) async fn update_view(
 		let ancestry_numbers = (min_number..=leaf_number).rev();
 		let mut ancestry_iter = ancestry_hashes.clone().zip(ancestry_numbers).peekable();
 		if let Some((hash, number)) = ancestry_iter.next() {
-			assert_matches!(
-				overseer_recv_with_timeout(virtual_overseer, Duration::from_millis(50)).await.unwrap(),
-				AllMessages::ChainApi(ChainApiMessage::BlockHeader(.., tx)) => {
-					let header = Header {
-						parent_hash: get_parent_hash(hash),
-						number,
-						state_root: Hash::zero(),
-						extrinsics_root: Hash::zero(),
-						digest: Default::default(),
-					};
-
-					tx.send(Ok(Some(header))).unwrap();
-				}
-			);
-
+			// fetch_ancestors is called first, which requests session for the leaf
 			assert_matches!(
 				overseer_recv_with_timeout(virtual_overseer, Duration::from_millis(50)).await.unwrap(),
 				AllMessages::RuntimeApi(
@@ -135,20 +121,37 @@ pub(super) async fn update_view(
 					tx.send(Ok(hashes)).unwrap();
 				}
 			);
-		}
 
-		for _ in ancestry_iter.clone() {
+			// fetch_ancestors checks session for each ancestor
+			for _ in ancestry_iter.clone() {
+				assert_matches!(
+					overseer_recv_with_timeout(virtual_overseer, Duration::from_millis(50)).await.unwrap(),
+					AllMessages::RuntimeApi(
+						RuntimeApiMessage::Request(
+							..,
+							RuntimeApiRequest::SessionIndexForChild(
+								tx
+							)
+						)
+					) => {
+						tx.send(Ok(1)).unwrap();
+					}
+				);
+			}
+
+			// Now fetch_fresh_leaf_and_insert_ancestry requests block headers
 			assert_matches!(
 				overseer_recv_with_timeout(virtual_overseer, Duration::from_millis(50)).await.unwrap(),
-				AllMessages::RuntimeApi(
-					RuntimeApiMessage::Request(
-						..,
-						RuntimeApiRequest::SessionIndexForChild(
-							tx
-						)
-					)
-				) => {
-					tx.send(Ok(1)).unwrap();
+				AllMessages::ChainApi(ChainApiMessage::BlockHeader(.., tx)) => {
+					let header = Header {
+						parent_hash: get_parent_hash(hash),
+						number,
+						state_root: Hash::zero(),
+						extrinsics_root: Hash::zero(),
+						digest: Default::default(),
+					};
+
+					tx.send(Ok(Some(header))).unwrap();
 				}
 			);
 		}

polkadot/node/network/collator-protocol/src/validator_side/tests/prospective_parachains.rs

@@ -127,11 +127,60 @@ pub(super) async fn update_view(
 		)
 		.await;
 
-		let min_number = leaf_number.saturating_sub(test_state.scheduling_lookahead);
+		// activate_leaf calls fetch_ancestors
+		assert_matches!(
+			overseer_recv(virtual_overseer).await,
+			AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+				_,
+				RuntimeApiRequest::SessionIndexForChild(tx)
+			)) => {
+				tx.send(Ok(test_state.session_index)).unwrap();
+			}
+		);
 
+		assert_matches!(
+			overseer_recv(virtual_overseer).await,
+			AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+				_,
+				RuntimeApiRequest::SchedulingLookahead(_, tx)
+			)) => {
+				tx.send(Ok(test_state.scheduling_lookahead)).unwrap();
+			}
+		);
+
+		let min_number = leaf_number.saturating_sub(test_state.scheduling_lookahead);
 		let ancestry_len = leaf_number + 1 - min_number;
 		let ancestry_hashes = std::iter::successors(Some(leaf_hash), |h| Some(get_parent_hash(*h)))
 			.take(ancestry_len as usize);
+
+		let returned_ancestors = assert_matches!(
+			overseer_recv(virtual_overseer).await,
+			AllMessages::ChainApi(ChainApiMessage::Ancestors {
+				k,
+				response_channel: tx,
+				..
+			}) => {
+				assert_eq!(k, test_state.scheduling_lookahead.saturating_sub(1) as usize);
+				let hashes: Vec<_> = ancestry_hashes.clone().skip(1).collect();
+				let returned = hashes.clone();
+				tx.send(Ok(hashes)).unwrap();
+				returned
+			}
+		);
+
+		// fetch_ancestors checks session for each ancestor that was returned
+		for _ in 0..returned_ancestors.len() {
+			assert_matches!(
+				overseer_recv(virtual_overseer).await,
+				AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+					_,
+					RuntimeApiRequest::SessionIndexForChild(tx)
+				)) => {
+					tx.send(Ok(test_state.session_index)).unwrap();
+				}
+			);
+		}
+
 		let ancestry_numbers = (min_number..=leaf_number).rev();
 		let ancestry_iter = ancestry_hashes.clone().zip(ancestry_numbers).peekable();
 
@@ -150,7 +199,18 @@ pub(super) async fn update_view(
 
 				let msg = match next_overseer_message.take() {
 					Some(msg) => msg,
-					None => overseer_recv(virtual_overseer).await,
+					None => match overseer_recv_with_timeout(
+						virtual_overseer,
+						Duration::from_millis(50),
+					)
+					.await
+					{
+						Some(msg) => msg,
+						None => {
+							// No message arrived - ancestry is cached
+							break
+						},
+					},
 				};
 
 				if !matches!(&msg, AllMessages::ChainApi(ChainApiMessage::BlockHeader(..))) {

polkadot/node/network/statement-distribution/src/v2/tests/mod.rs

@@ -600,6 +600,68 @@ async fn handle_leaf_activation(
 		claim_queue,
 	} = leaf;
 
+	// activate_leaf calls fetch_ancestors first
+	assert_matches!(
+		virtual_overseer.recv().await,
+		AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+			_,
+			RuntimeApiRequest::SessionIndexForChild(tx),
+		)) => {
+			tx.send(Ok(*session)).unwrap();
+		}
+	);
+
+	assert_matches!(
+		virtual_overseer.recv().await,
+		AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+			_,
+			RuntimeApiRequest::SchedulingLookahead(_, tx),
+		)) => {
+			// Assuming scheduling lookahead of 2 for tests
+			tx.send(Ok(2)).unwrap();
+		}
+	);
+
+	let ancestors = assert_matches!(
+		virtual_overseer.recv().await,
+		AllMessages::ChainApi(ChainApiMessage::Ancestors {
+			k,
+			response_channel: tx,
+			..
+		}) => {
+			// Calculate ancestors based on block number
+			let mut ancestors = Vec::new();
+			let mut current_hash = *parent_hash;
+			let mut current_number = *number - 1;
+
+			for _ in 0..k {
+				if current_number == 0 {
+					break;
+				}
+				ancestors.push(current_hash);
+				// For tests, generate predictable parent hashes
+				current_hash = Hash::repeat_byte(current_hash.as_ref()[0].wrapping_sub(1));
+				current_number -= 1;
+			}
+
+			tx.send(Ok(ancestors.clone())).unwrap();
+			ancestors
+		}
+	);
+
+	// fetch_ancestors checks session for each returned ancestor
+	for _ in 0..ancestors.len() {
+		assert_matches!(
+			virtual_overseer.recv().await,
+			AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+				_,
+				RuntimeApiRequest::SessionIndexForChild(tx),
+			)) => {
+				tx.send(Ok(*session)).unwrap();
+			}
+		);
+	}
+
 	let header = Header {
 		parent_hash: *parent_hash,
 		number: *number,