From 0e305bd7ddee6e5b4b12623cf4fb0e16152cb0ef Mon Sep 17 00:00:00 2001
From: Greg Fitzgerald <greg@solana.com>
Date: Fri, 20 Jul 2018 19:05:39 -0400
Subject: [PATCH] Add poll_for_signature

---
 src/bin/client-demo.rs |  0
 src/drone.rs           |  9 ++++-----
 src/thin_client.rs     | 23 +++++++++++++++++++----
 3 files changed, 23 insertions(+), 9 deletions(-)
 mode change 100755 => 100644 src/bin/client-demo.rs

diff --git a/src/bin/client-demo.rs b/src/bin/client-demo.rs
old mode 100755
new mode 100644
diff --git a/src/drone.rs b/src/drone.rs
index 524ff18ee..de84cbdac 100644
--- a/src/drone.rs
+++ b/src/drone.rs
@@ -96,7 +96,6 @@ impl Drone {
     }
 
     pub fn send_airdrop(&mut self, req: DroneRequest) -> Result<Signature, io::Error> {
-        let tx: Transaction;
         let request_amount: u64;
         let requests_socket = UdpSocket::bind("0.0.0.0:0").unwrap();
         let transactions_socket = UdpSocket::bind("0.0.0.0:0").unwrap();
@@ -109,7 +108,7 @@ impl Drone {
         );
         let last_id = client.get_last_id();
 
-        match req {
+        let tx = match req {
             DroneRequest::GetAirdrop {
                 airdrop_request_amount,
                 client_public_key,
@@ -119,14 +118,14 @@ impl Drone {
                     airdrop_request_amount, client_public_key
                 );
                 request_amount = airdrop_request_amount;
-                tx = Transaction::new(
+                Transaction::new(
                     &self.mint_keypair,
                     client_public_key,
                     airdrop_request_amount as i64,
                     last_id,
-                );
+                )
             }
-        }
+        };
         if self.check_request_limit(request_amount) {
             self.request_current += request_amount;
             metrics::submit(
diff --git a/src/thin_client.rs b/src/thin_client.rs
index 64bb251cc..c74494c47 100644
--- a/src/thin_client.rs
+++ b/src/thin_client.rs
@@ -217,6 +217,19 @@ impl ThinClient {
         balance
     }
 
+    /// Poll the server to confirm a transaction.
+    pub fn poll_for_signature(&mut self, sig: &Signature) -> io::Result<()> {
+        let now = Instant::now();
+        while !self.check_signature(sig) {
+            if now.elapsed().as_secs() > 1 {
+                // TODO: Return a better error.
+                return Err(io::Error::new(io::ErrorKind::Other, "signature not found"));
+            }
+            sleep(Duration::from_millis(100));
+        }
+        Ok(())
+    }
+
     /// Check a signature in the bank. This method blocks
     /// until the server sends a response.
     pub fn check_signature(&mut self, sig: &Signature) -> bool {
@@ -304,10 +317,11 @@ mod tests {
             transactions_socket,
         );
         let last_id = client.get_last_id();
-        let _sig = client
+        let sig = client
             .transfer(500, &alice.keypair(), bob_pubkey, &last_id)
             .unwrap();
-        let balance = client.poll_get_balance(&bob_pubkey);
+        client.poll_for_signature(&sig).unwrap();
+        let balance = client.get_balance(&bob_pubkey);
         assert_eq!(balance.unwrap(), 500);
         exit.store(true, Ordering::Relaxed);
         server.join().unwrap();
@@ -361,9 +375,10 @@ mod tests {
             contract.tokens = 502;
             contract.plan = Plan::Budget(Budget::new_payment(502, bob_pubkey));
         }
-        let _sig = client.transfer_signed(&tr2).unwrap();
+        let sig = client.transfer_signed(&tr2).unwrap();
+        client.poll_for_signature(&sig).unwrap();
 
-        let balance = client.poll_get_balance(&bob_pubkey);
+        let balance = client.get_balance(&bob_pubkey);
         assert_eq!(balance.unwrap(), 500);
         exit.store(true, Ordering::Relaxed);
         server.join().unwrap();