add a couple more tests

polkadot/network/src/collator_pool.rs

@@ -211,8 +211,37 @@ mod tests {
 		let relay_parent = [1; 32].into();
 
 		assert_eq!(pool.on_new_collator(primary, para_id.clone()), Role::Primary);
-		let (tx, rx) = oneshot::channel();
-		pool.await_collation(relay_parent, para_id, tx);
+		let (tx1, rx1) = oneshot::channel();
+		let (tx2, rx2) = oneshot::channel();
+		pool.await_collation(relay_parent, para_id, tx1);
+		pool.await_collation(relay_parent, para_id, tx2);
+		pool.on_collation(primary, relay_parent, Collation {
+			receipt: CandidateReceipt {
+				parachain_index: para_id,
+				collator: primary.into(),
+				signature: H512::from([2; 64]).into(),
+				head_data: HeadData(vec![1, 2, 3]),
+				balance_uploads: vec![],
+				egress_queue_roots: vec![],
+				fees: 0,
+				block_data_hash: [3; 32].into(),
+			},
+			block_data: BlockData(vec![4, 5, 6]),
+		});
+
+		rx1.wait().unwrap();
+		rx2.wait().unwrap();
+	}
+
+	#[test]
+	fn collate_before_await() {
+		let mut pool = CollatorPool::new();
+		let para_id: ParaId = 5.into();
+		let primary = [0; 32].into();
+		let relay_parent = [1; 32].into();
+
+		assert_eq!(pool.on_new_collator(primary, para_id.clone()), Role::Primary);
+
 		pool.on_collation(primary, relay_parent, Collation {
 			receipt: CandidateReceipt {
 				parachain_index: para_id,
@@ -227,6 +256,8 @@ mod tests {
 			block_data: BlockData(vec![4, 5, 6]),
 		});
 
+		let (tx, rx) = oneshot::channel();
+		pool.await_collation(relay_parent, para_id, tx);
 		rx.wait().unwrap();
 	}
 }

polkadot/network/src/lib.rs

@@ -422,15 +422,16 @@ impl Specialization<Block> for PolkadotProtocol {
 		}
 
 		let validator = status.roles.iter().any(|r| *r == message::Role::Authority);
+		let send_key = validator || local_status.collating_for.is_some();
+
 		self.peers.insert(peer_id, PeerInfo {
 			status: local_status,
 			session_keys: Default::default(),
 			validator,
 		});
 
 		self.consensus_gossip.new_peer(ctx, peer_id, &status.roles);
-
-		if let (true, &Some(ref consensus)) = (validator, &self.live_consensus) {
+		if let (true, &Some(ref consensus)) = (send_key, &self.live_consensus) {
 			send_polkadot_message(
 				ctx,
 				peer_id,

polkadot/network/src/tests.rs

@@ -110,11 +110,12 @@ fn sends_session_key() {
 	let parent_hash = [0; 32].into();
 	let local_key = [1; 32].into();
 
-	let status = Status { collating_for: None };
+	let validator_status = Status { collating_for: None };
+	let collator_status = Status { collating_for: Some(([2; 32].into(), 5.into())) };
 
 	{
 		let mut ctx = TestContext::default();
-		protocol.on_connect(&mut ctx, peer_a, make_status(&status, vec![Role::Authority]));
+		protocol.on_connect(&mut ctx, peer_a, make_status(&validator_status, vec![Role::Authority]));
 		assert!(ctx.messages.is_empty());
 	}
 
@@ -128,7 +129,7 @@ fn sends_session_key() {
 
 	{
 		let mut ctx = TestContext::default();
-		protocol.on_connect(&mut ctx, peer_b, make_status(&status, vec![Role::Authority]));
+		protocol.on_connect(&mut ctx, peer_b, make_status(&collator_status, vec![]));
 		assert!(ctx.has_message(peer_b, Message::SessionKey(parent_hash, local_key)));
 	}
 }
@@ -207,3 +208,24 @@ fn fetches_from_those_with_knowledge() {
 		assert_eq!(recv.wait().unwrap(), block_data);
 	}
 }
+
+#[test]
+fn remove_bad_collator() {
+	let mut protocol = PolkadotProtocol::new();
+
+	let peer_id = 1;
+	let account_id = [2; 32].into();
+
+	let status = Status { collating_for: Some((account_id, 5.into())) };
+
+	{
+		let mut ctx = TestContext::default();
+		protocol.on_connect(&mut ctx, peer_id, make_status(&status, vec![]));
+	}
+
+	{
+		let mut ctx = TestContext::default();
+		protocol.disconnect_bad_collator(&mut ctx, account_id);
+		assert!(ctx.disabled.contains(&peer_id));
+	}
+}