diff --git a/secret_store/src/key_server.rs b/secret_store/src/key_server.rs
index 969782ca2..c83e460f3 100644
--- a/secret_store/src/key_server.rs
+++ b/secret_store/src/key_server.rs
@@ -192,10 +192,13 @@ impl Drop for KeyServerCore {
 pub mod tests {
 	use std::time;
 	use std::sync::Arc;
+	use std::net::SocketAddr;
+	use std::collections::BTreeMap;
 	use ethcrypto;
 	use ethkey::{self, Secret, Random, Generator};
 	use acl_storage::tests::DummyAclStorage;
 	use key_storage::tests::DummyKeyStorage;
+	use key_server_set::tests::MapKeyServerSet;
 	use key_server_cluster::math;
 	use util::H256;
 	use types::all::{Error, Public, ClusterConfiguration, NodeAddress, RequestSignature, ServerKeyId,
@@ -253,8 +256,11 @@ pub mod tests {
 					})).collect(),
 				allow_connecting_to_higher_nodes: false,
 			}).collect();
+		let key_servers_set: BTreeMap<Public, SocketAddr> = configs[0].nodes.iter()
+			.map(|(k, a)| (k.clone(), format!("{}:{}", a.address, a.port).parse().unwrap()))
+			.collect();
 		let key_servers: Vec<_> = configs.into_iter().map(|cfg|
-			KeyServerImpl::new(&cfg, Arc::new(DummyAclStorage::default()), Arc::new(DummyKeyStorage::default())).unwrap()
+			KeyServerImpl::new(&cfg, Arc::new(MapKeyServerSet::new(key_servers_set.clone())), Arc::new(DummyAclStorage::default()), Arc::new(DummyKeyStorage::default())).unwrap()
 		).collect();
 
 		// wait until connections are established. It is fast => do not bother with events here
diff --git a/secret_store/src/key_server_cluster/cluster.rs b/secret_store/src/key_server_cluster/cluster.rs
index 649ad4bc9..d77a82431 100644
--- a/secret_store/src/key_server_cluster/cluster.rs
+++ b/secret_store/src/key_server_cluster/cluster.rs
@@ -989,7 +989,7 @@ pub mod tests {
 	use parking_lot::Mutex;
 	use tokio_core::reactor::Core;
 	use ethkey::{Random, Generator, Public};
-	use key_server_cluster::{NodeId, SessionId, Error, DummyAclStorage, DummyKeyStorage};
+	use key_server_cluster::{NodeId, SessionId, Error, DummyAclStorage, DummyKeyStorage, MapKeyServerSet};
 	use key_server_cluster::message::Message;
 	use key_server_cluster::cluster::{Cluster, ClusterCore, ClusterConfiguration};
 	use key_server_cluster::generation_session::{Session as GenerationSession, SessionState as GenerationSessionState};
@@ -1059,7 +1059,7 @@ pub mod tests {
 	}
 
 	pub fn all_connections_established(cluster: &Arc<ClusterCore>) -> bool {
-		cluster.config().nodes.keys()
+		cluster.config().key_server_set.get().keys()
 			.filter(|p| *p != cluster.config().self_key_pair.public())
 			.all(|p| cluster.connection(p).is_some())
 	}
@@ -1070,9 +1070,9 @@ pub mod tests {
 			threads: 1,
 			self_key_pair: key_pairs[i].clone(),
 			listen_address: ("127.0.0.1".to_owned(), ports_begin + i as u16),
-			nodes: key_pairs.iter().enumerate()
-				.map(|(j, kp)| (kp.public().clone(), ("127.0.0.1".into(), ports_begin + j as u16)))
-				.collect(),
+			key_server_set: Arc::new(MapKeyServerSet::new(key_pairs.iter().enumerate()
+				.map(|(j, kp)| (kp.public().clone(), format!("127.0.0.1:{}", ports_begin + j as u16).parse().unwrap()))
+				.collect())),
 			allow_connecting_to_higher_nodes: false,
 			key_storage: Arc::new(DummyKeyStorage::default()),
 			acl_storage: Arc::new(DummyAclStorage::default()),
diff --git a/secret_store/src/key_server_cluster/mod.rs b/secret_store/src/key_server_cluster/mod.rs
index 11c32d528..8f6ae4add 100644
--- a/secret_store/src/key_server_cluster/mod.rs
+++ b/secret_store/src/key_server_cluster/mod.rs
@@ -34,6 +34,8 @@ pub use self::decryption_session::Session as DecryptionSession;
 pub use super::key_storage::tests::DummyKeyStorage;
 #[cfg(test)]
 pub use super::acl_storage::tests::DummyAclStorage;
+#[cfg(test)]
+pub use super::key_server_set::tests::MapKeyServerSet;
 
 pub type SessionId = ServerKeyId;
 
diff --git a/secret_store/src/key_server_set.rs b/secret_store/src/key_server_set.rs
index 47f033db1..e17dceed5 100644
--- a/secret_store/src/key_server_set.rs
+++ b/secret_store/src/key_server_set.rs
@@ -175,3 +175,30 @@ impl CachedContract {
 		self.contract_addr = new_contract_address;
 	}
 }
+
+#[cfg(test)]
+pub mod tests {
+	use std::collections::BTreeMap;
+	use std::net::SocketAddr;
+	use ethkey::Public;
+	use super::KeyServerSet;
+
+	#[derive(Default)]
+	pub struct MapKeyServerSet {
+		nodes: BTreeMap<Public, SocketAddr>,
+	}
+
+	impl MapKeyServerSet {
+		pub fn new(nodes: BTreeMap<Public, SocketAddr>) -> Self {
+			MapKeyServerSet {
+				nodes: nodes,
+			}
+		}
+	}
+
+	impl KeyServerSet for MapKeyServerSet {
+		fn get(&self) -> BTreeMap<Public, SocketAddr> {
+			self.nodes.clone()
+		}
+	}
+}