diff --git a/state-chain/pallets/cf-elections/src/electoral_systems/mocks.rs b/state-chain/pallets/cf-elections/src/electoral_systems/mocks.rs
index 4a5657f0334..b199f88f9ed 100644
--- a/state-chain/pallets/cf-elections/src/electoral_systems/mocks.rs
+++ b/state-chain/pallets/cf-elections/src/electoral_systems/mocks.rs
@@ -7,7 +7,6 @@ use crate::{
 	},
 	UniqueMonotonicIdentifier,
 };
-use cf_traits::mocks::time_source::Mock;
 use frame_support::{CloneNoBound, DebugNoBound, EqNoBound, PartialEqNoBound};
 
 pub mod access;
diff --git a/state-chain/pallets/cf-elections/src/electoral_systems/tests/unsafe_median.rs b/state-chain/pallets/cf-elections/src/electoral_systems/tests/unsafe_median.rs
index 4e0b2311e66..dff2e35e319 100644
--- a/state-chain/pallets/cf-elections/src/electoral_systems/tests/unsafe_median.rs
+++ b/state-chain/pallets/cf-elections/src/electoral_systems/tests/unsafe_median.rs
@@ -2,7 +2,7 @@ use cf_primitives::AuthorityCount;
 
 use super::{mocks::*, register_checks};
 use crate::{
-	electoral_system::{ConsensusStatus, ElectoralReadAccess},
+	electoral_system::ConsensusStatus,
 	electoral_systems::{tests::utils::generate_votes, unsafe_median::*},
 };
 
@@ -19,66 +19,65 @@ fn with_default_context() -> TestContext<SimpleUnsafeMedian> {
 	with_default_setup().build_with_initial_election()
 }
 
-// TODO: Remove second arg from macro
-// register_checks! {
-// 	SimpleUnsafeMedian {
-// 		started_at_initial_state(_post, _a) {
-// 			assert_eq!(
-// 				MockAccess::<SimpleUnsafeMedian>::unsynchronised_state().unwrap(),
-// 				INIT_UNSYNCHRONISED_STATE,
-// 				"Expected initial state pre-finalization."
-// 			);
-// 		},
-// 		ended_at_initial_state(_pre, _a) {
-// 			assert_eq!(
-// 				MockAccess::<SimpleUnsafeMedian>::unsynchronised_state().unwrap(),
-// 				INIT_UNSYNCHRONISED_STATE,
-// 				"Expected initial state post-finalization."
-// 			);
-// 		},
-// 		ended_at_new_state(_pre, _a) {
-// 			assert_eq!(
-// 				MockAccess::<SimpleUnsafeMedian>::unsynchronised_state().unwrap(),
-// 				NEW_UNSYNCHRONISED_STATE,
-// 				"Expected new state post-finalization."
-// 			);
-// 		},
-// 	}
-// }
+register_checks! {
+	SimpleUnsafeMedian {
+		started_at_initial_state(pre_finalize, _a) {
+			assert_eq!(
+				pre_finalize.unsynchronised_state,
+				INIT_UNSYNCHRONISED_STATE,
+				"Expected initial state pre-finalization."
+			);
+		},
+		ended_at_initial_state(_pre, post_finalize) {
+			assert_eq!(
+				post_finalize.unsynchronised_state,
+				INIT_UNSYNCHRONISED_STATE,
+				"Expected initial state post-finalization."
+			);
+		},
+		ended_at_new_state(_pre, post_finalize) {
+			assert_eq!(
+				post_finalize.unsynchronised_state,
+				NEW_UNSYNCHRONISED_STATE,
+				"Expected new state post-finalization."
+			);
+		},
+	}
+}
 
-// #[test]
-// fn if_consensus_update_unsynchronised_state() {
-// 	with_default_context()
-// 		.force_consensus_update(ConsensusStatus::Gained {
-// 			most_recent: None,
-// 			new: NEW_UNSYNCHRONISED_STATE,
-// 		})
-// 		.test_on_finalize(
-// 			&(),
-// 			|_| {},
-// 			vec![
-// 				Check::started_at_initial_state(),
-// 				Check::ended_at_new_state(),
-// 				Check::last_election_deleted(),
-// 				Check::election_id_incremented(),
-// 			],
-// 		);
-// }
+#[test]
+fn if_consensus_update_unsynchronised_state() {
+	with_default_context()
+		.force_consensus_update(ConsensusStatus::Gained {
+			most_recent: None,
+			new: NEW_UNSYNCHRONISED_STATE,
+		})
+		.test_on_finalize(
+			&(),
+			|_| {},
+			vec![
+				Check::started_at_initial_state(),
+				Check::ended_at_new_state(),
+				Check::last_election_deleted(),
+				Check::election_id_incremented(),
+			],
+		);
+}
 
-// #[test]
-// fn if_no_consensus_do_not_update_unsynchronised_state() {
-// 	with_default_context()
-// 		.force_consensus_update(ConsensusStatus::None)
-// 		.test_on_finalize(
-// 			&(),
-// 			|_| {},
-// 			vec![
-// 				Check::started_at_initial_state(),
-// 				Check::ended_at_initial_state(),
-// 				Check::assert_unchanged(),
-// 			],
-// 		);
-// }
+#[test]
+fn if_no_consensus_do_not_update_unsynchronised_state() {
+	with_default_context()
+		.force_consensus_update(ConsensusStatus::None)
+		.test_on_finalize(
+			&(),
+			|_| {},
+			vec![
+				Check::started_at_initial_state(),
+				Check::ended_at_initial_state(),
+				Check::assert_unchanged(),
+			],
+		);
+}
 
 #[test]
 fn check_consensus_correctly_calculates_median_when_all_authorities_vote() {