From 6dc4633bddfac1cf43a1024444ea09f0302d6946 Mon Sep 17 00:00:00 2001
From: ZHANG Dapeng <zdapeng@google.com>
Date: Mon, 28 Nov 2016 18:13:22 -0800
Subject: [PATCH] stub: add withWaitForReady

---
 .../main/java/io/grpc/stub/AbstractStub.java  |  8 ++++++
 .../java/io/grpc/stub/AbstractStubTest.java   | 26 +++++++++++++++++--
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/stub/src/main/java/io/grpc/stub/AbstractStub.java b/stub/src/main/java/io/grpc/stub/AbstractStub.java
index 802f1b307b..87f0e56482 100644
--- a/stub/src/main/java/io/grpc/stub/AbstractStub.java
+++ b/stub/src/main/java/io/grpc/stub/AbstractStub.java
@@ -181,4 +181,12 @@ public abstract class AbstractStub<S extends AbstractStub<S>> {
   public final S withCallCredentials(CallCredentials credentials) {
     return build(channel, callOptions.withCallCredentials(credentials));
   }
+
+  /**
+   * Returns a new stub that uses the 'wait for ready' call option.
+   */
+  @ExperimentalApi("https://github.com/grpc/grpc-java/issues/1915")
+  public final S withWaitForReady() {
+    return build(channel, callOptions.withWaitForReady());
+  }
 }
diff --git a/stub/src/test/java/io/grpc/stub/AbstractStubTest.java b/stub/src/test/java/io/grpc/stub/AbstractStubTest.java
index a5c9e21399..fa7e790ebd 100644
--- a/stub/src/test/java/io/grpc/stub/AbstractStubTest.java
+++ b/stub/src/test/java/io/grpc/stub/AbstractStubTest.java
@@ -31,12 +31,18 @@
 
 package io.grpc.stub;
 
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
 import io.grpc.CallOptions;
 import io.grpc.Channel;
+
+import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 
 @RunWith(JUnit4.class)
 public class AbstractStubTest {
@@ -44,6 +50,11 @@ public class AbstractStubTest {
   @Mock
   Channel channel;
 
+  @Before
+  public void setup() {
+    MockitoAnnotations.initMocks(this);
+  }
+
   @Test(expected = NullPointerException.class)
   public void channelMustNotBeNull() {
     new NoopStub(null);
@@ -55,8 +66,19 @@ public class AbstractStubTest {
   }
 
   @Test(expected = NullPointerException.class)
-  public void callOptionsAndChannelMustNotBeNull() {
-    new NoopStub(null, null);
+  public void channelMustNotBeNull2() {
+    new NoopStub(null, CallOptions.DEFAULT);
+  }
+
+  @Test()
+  public void withWaitForReady() {
+    NoopStub stub = new NoopStub(channel);
+    CallOptions callOptions = stub.getCallOptions();
+    assertFalse(callOptions.isWaitForReady());
+
+    stub = stub.withWaitForReady();
+    callOptions = stub.getCallOptions();
+    assertTrue(callOptions.isWaitForReady());
   }
 
   class NoopStub extends AbstractStub<NoopStub> {