From 2898474439acb6c4bc1e6d05af64a4c7f1a303b6 Mon Sep 17 00:00:00 2001
From: Tony Arcieri <tony@iqlusion.io>
Date: Tue, 6 Apr 2021 13:11:08 -0700
Subject: [PATCH] rpc: add `Client::wait_until_healthy` method

Adds a method which repeatedly polls the `/health` endpoint until it
returns a successful response or a specified timeout elapses.

This is useful when writing tests that need to wait for a node to boot
up (e.g. inside of a docker container).
---
 rpc/src/client.rs | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/rpc/src/client.rs b/rpc/src/client.rs
index cb4070f55..bcab4946e 100644
--- a/rpc/src/client.rs
+++ b/rpc/src/client.rs
@@ -14,14 +14,17 @@ pub use transport::websocket::{WebSocketClient, WebSocketClientDriver, WebSocket
 
 use crate::endpoint::validators::DEFAULT_VALIDATORS_PER_PAGE;
 use crate::endpoint::*;
+use crate::error::Error;
 use crate::paging::Paging;
 use crate::query::Query;
 use crate::{Order, Result, SimpleRequest};
 use async_trait::async_trait;
+use std::time::Duration;
 use tendermint::abci::{self, Transaction};
 use tendermint::block::Height;
 use tendermint::evidence::Evidence;
 use tendermint::Genesis;
+use tokio::time;
 
 /// Provides lightweight access to the Tendermint RPC. It gives access to all
 /// endpoints with the exception of the event subscription-related ones.
@@ -221,6 +224,31 @@ pub trait Client {
             .await
     }
 
+    /// Poll the `/health` endpoint until it returns a successful result or
+    /// the given `timeout` has elapsed.
+    async fn wait_until_healthy<T>(&self, timeout: T) -> Result<()>
+    where
+        T: Into<Duration> + Send,
+    {
+        let timeout = timeout.into();
+        let poll_interval = Duration::from_millis(200);
+        let mut attempts_remaining = timeout.as_millis() / poll_interval.as_millis();
+
+        while self.health().await.is_err() {
+            if attempts_remaining == 0 {
+                return Err(Error::client_internal_error(format!(
+                    "timed out waiting for healthy response after {}ms",
+                    timeout.as_millis()
+                )));
+            }
+
+            attempts_remaining -= 1;
+            time::sleep(poll_interval).await;
+        }
+
+        Ok(())
+    }
+
     /// Perform a request against the RPC endpoint
     async fn perform<R>(&self, request: R) -> Result<R::Response>
     where