Fix beacon_chain_sim, nitpicks

tests/beacon_chain_sim/src/main.rs

@@ -1,5 +1,5 @@
 use node_test_rig::{
-    environment::{EnvironmentBuilder, RuntimeContext},
+    environment::{Environment, EnvironmentBuilder, RuntimeContext},
     testing_client_config, ClientConfig, ClientGenesis, LocalBeaconNode, LocalValidatorClient,
     ProductionClient, ValidatorConfig,
 };
@@ -64,10 +64,13 @@ fn simulation(num_nodes: usize, validators_per_node: usize) -> Result<(), String
                 .spec
                 .clone();
 
+            let context = env.service_context(format!("validator_{}", i));
+
             let indices =
                 (i * validators_per_node..(i + 1) * validators_per_node).collect::<Vec<_>>();
             new_validator_client(
-                env.service_context(format!("validator_{}", i)),
+                &mut env,
+                context,
                 node,
                 ValidatorConfig::default(),
                 &indices,
@@ -100,7 +103,10 @@ fn new_with_bootnode_via_enr<E: EthSpec>(
     BeaconNode::production(context, config)
 }
 
+// Note: this function will block until the validator can connect to the beaco node. It is
+// recommended to ensure that the beacon node is running first.
 fn new_validator_client<E: EthSpec>(
+    env: &mut Environment<E>,
     context: RuntimeContext<E>,
     beacon_node: &BeaconNode<E>,
     base_config: ValidatorConfig,
@@ -115,5 +121,11 @@ fn new_validator_client<E: EthSpec>(
 
     config.http_server = format!("http://{}:{}", socket_addr.ip(), socket_addr.port());
 
-    LocalValidatorClient::production_with_insecure_keypairs(context, config, keypair_indices)
+    env.runtime()
+        .block_on(LocalValidatorClient::production_with_insecure_keypairs(
+            context,
+            config,
+            keypair_indices,
+        ))
+        .expect("should start validator")
 }

tests/node_test_rig/src/lib.rs

@@ -106,7 +106,7 @@ impl<E: EthSpec> LocalValidatorClient<E> {
         context: RuntimeContext<E>,
         mut config: ValidatorConfig,
         keypair_indices: &[usize],
-    ) -> Self {
+    ) -> impl Future<Item = Self, Error = String> {
         // Creates a temporary directory that will be deleted once this `TempDir` is dropped.
         let datadir = TempDir::new("lighthouse-beacon-node")
             .expect("should create temp directory for client datadir");
@@ -120,7 +120,10 @@ impl<E: EthSpec> LocalValidatorClient<E> {
     ///
     /// - The validator created is using the same types as the node we use in production.
     /// - It is recommended to use `production_with_insecure_keypairs` for testing.
-    pub fn production(context: RuntimeContext<E>, config: ValidatorConfig) -> Self {
+    pub fn production(
+        context: RuntimeContext<E>,
+        config: ValidatorConfig,
+    ) -> impl Future<Item = Self, Error = String> {
         // Creates a temporary directory that will be deleted once this `TempDir` is dropped.
         let datadir = TempDir::new("lighthouse-validator")
             .expect("should create temp directory for client datadir");
@@ -128,17 +131,18 @@ impl<E: EthSpec> LocalValidatorClient<E> {
         Self::new(context, config, datadir)
     }
 
-    fn new(context: RuntimeContext<E>, mut config: ValidatorConfig, datadir: TempDir) -> Self {
+    fn new(
+        context: RuntimeContext<E>,
+        mut config: ValidatorConfig,
+        datadir: TempDir,
+    ) -> impl Future<Item = Self, Error = String> {
         config.data_dir = datadir.path().into();
 
-        let client = ProductionValidatorClient::new(context, config)
-            .wait()
-            .expect("should start validator client");
-
-        client
-            .start_service()
-            .expect("should start validator client");
-
-        Self { client, datadir }
+        ProductionValidatorClient::new(context, config).map(move |mut client| {
+            client
+                .start_service()
+                .expect("should start validator services");
+            Self { client, datadir }
+        })
     }
 }