Use Rpc.await() functionality.

1 files changed, 44 insertions, 25 deletions

diff --git a/same/src/main/java/com/orbekk/same/DirectoryApp.java b/same/src/main/java/com/orbekk/same/DirectoryApp.java
index e75b10f..54636bc 100644
--- a/same/src/main/java/com/orbekk/same/DirectoryApp.java
+++ b/same/src/main/java/com/orbekk/same/DirectoryApp.java
@@ -11,6 +11,9 @@ import com.google.protobuf.RpcCallback;
 import com.orbekk.protobuf.NewRpcChannel;
 import com.orbekk.protobuf.Rpc;
 import com.orbekk.protobuf.RpcChannel;
+import com.orbekk.same.Services.Directory;
+import com.orbekk.same.Services.Empty;
+import com.orbekk.same.Services.MasterState;
 import com.orbekk.same.Services.NetworkDirectory;
 
 public class DirectoryApp {
@@ -24,24 +27,26 @@ public class DirectoryApp {
         this.args = args;
     }
     
-    public void run() {
-        String host = args[0];
-        int port = Integer.valueOf(args[1]);
-        NewRpcChannel channel = null;
-        try {
-            channel = NewRpcChannel.create(host, port);
-        } catch (UnknownHostException e1) {
-            e1.printStackTrace();
-            System.exit(1);
-        } catch (IOException e1) {
-            e1.printStackTrace();
-            System.exit(1);
+    private void addTestNetwork(Directory directory) throws InterruptedException {
+        final Rpc rpc = new Rpc();
+        MasterState request = MasterState.newBuilder()
+                .setNetworkName("Test network")
+                .setMasterUrl("invalid:invalid")
+                .build();
+        RpcCallback<Empty> done = new RpcCallback<Empty>() {
+            @Override public void run(Empty unused) {
+            }
+        };
+        directory.registerNetwork(rpc, request, done);
+        rpc.await();
+        if (rpc.isOk()) {
+            System.out.println("Added network.");
         }
-        Services.Directory directory = Services.Directory.newStub(channel);
-
-        final CountDownLatch finished = new CountDownLatch(1);
+    }
+    
+    private void listNetworks(Directory directory) throws InterruptedException {
         final Rpc rpc = new Rpc();
-        RpcCallback<NetworkDirectory> callback =
+        RpcCallback<NetworkDirectory> done =
                 new RpcCallback<NetworkDirectory>() {
             @Override public void run(NetworkDirectory directory) {
                 if (rpc.failed()) {
@@ -54,22 +59,36 @@ public class DirectoryApp {
                                 network.getMasterUrl());
                     }
                 }
-                finished.countDown();
             }
         };
-        directory.getNetworks(rpc, Services.Empty.getDefaultInstance(),
-                callback);
+        directory.getNetworks(rpc, Empty.getDefaultInstance(), done);
+        rpc.await();
+    }
+    
+    public void run() throws InterruptedException {
+        String host = args[0];
+        int port = Integer.valueOf(args[1]);
+        NewRpcChannel channel = null;
         try {
-            finished.await();
-        } catch (InterruptedException e) {
-            e.printStackTrace();
+            channel = NewRpcChannel.create(host, port);
+        } catch (UnknownHostException e1) {
+            e1.printStackTrace();
+            System.exit(1);
+        } catch (IOException e1) {
+            e1.printStackTrace();
+            System.exit(1);
         }
-        
-        System.out.println("Closing channel.");
+        Services.Directory directory = Services.Directory.newStub(channel);
+        addTestNetwork(directory);
+        listNetworks(directory);
         channel.close();
     }
     
     public static void main(String[] args) {
-        new DirectoryApp(args).run();
+        try {
+            new DirectoryApp(args).run();
+        } catch (InterruptedException e) {
+            throw new AssertionError("Should not be interrupted.");
+        }
     }
 }