Improve tests for BABE

core/consensus/babe/primitives/src/digest.rs

@@ -246,6 +246,10 @@ impl<Hash> CompatibleDigestItem for DigestItem<Hash> where
 
 	fn as_next_epoch_descriptor(&self) -> Option<NextEpochDescriptor> {
 		self.try_to(OpaqueDigestItemId::Consensus(&BABE_ENGINE_ID))
+			.and_then(|x: super::ConsensusLog| match x {
+				super::ConsensusLog::NextEpochData(n) => Some(n),
+				_ => None,
+			})
 	}
 }
 

core/consensus/babe/primitives/src/lib.rs

@@ -105,6 +105,12 @@ impl Epoch {
 			randomness: descriptor.randomness,
 		}
 	}
+
+	/// Produce the "end slot" of the epoch. This is NOT inclusive to the epoch,
+	// i.e. the slots covered by the epoch are `self.start_slot .. self.end_slot()`.
+	pub fn end_slot(&self) -> SlotNumber {
+		self.start_slot + self.duration
+	}
 }
 
 /// An consensus log item for BABE.

core/consensus/babe/src/aux_schema.rs

@@ -24,7 +24,7 @@ use client::error::{Result as ClientResult, Error as ClientError};
 use sr_primitives::traits::Block as BlockT;
 use babe_primitives::BabeBlockWeight;
 
-use super::{EpochChanges, SharedEpochChanges};
+use super::{epoch_changes::EpochChangesFor, SharedEpochChanges};
 
 const BABE_EPOCH_CHANGES: &[u8] = b"babe_epoch_changes";
 
@@ -50,7 +50,7 @@ fn load_decode<B, T>(backend: &B, key: &[u8]) -> ClientResult<Option<T>>
 pub(crate) fn load_epoch_changes<Block: BlockT, B: AuxStore>(
 	backend: &B,
 ) -> ClientResult<SharedEpochChanges<Block>> {
-	let epoch_changes = load_decode::<_, EpochChanges<Block>>(backend, BABE_EPOCH_CHANGES)?
+	let epoch_changes = load_decode::<_, EpochChangesFor<Block>>(backend, BABE_EPOCH_CHANGES)?
 		.map(Into::into)
 		.unwrap_or_else(|| {
 			info!(target: "babe",
@@ -64,7 +64,7 @@ pub(crate) fn load_epoch_changes<Block: BlockT, B: AuxStore>(
 
 /// Update the epoch changes on disk after a change.
 pub(crate) fn write_epoch_changes<Block: BlockT, F, R>(
-	epoch_changes: &EpochChanges<Block>,
+	epoch_changes: &EpochChangesFor<Block>,
 	write_aux: F,
 ) -> R where
 	F: FnOnce(&[(&'static [u8], &[u8])]) -> R,