s2a: Improve S2AStubTest clarity and performance

s2a/src/main/java/io/grpc/s2a/internal/handshaker/S2AStub.java

@@ -43,6 +43,7 @@ public class S2AStub implements AutoCloseable {
   private final BlockingQueue<Result> responses = new ArrayBlockingQueue<>(10);
   private S2AServiceGrpc.S2AServiceStub serviceStub;
   private StreamObserver<SessionReq> writer;
+  private long deadlineSeconds = HANDSHAKE_RPC_DEADLINE_SECS;
   private boolean doneReading = false;
   private boolean doneWriting = false;
   private boolean isClosed = false;
@@ -53,6 +54,14 @@ public static S2AStub newInstance(S2AServiceGrpc.S2AServiceStub serviceStub) {
     return new S2AStub(serviceStub);
   }
 
+  @VisibleForTesting
+  static S2AStub newInstanceWithDeadline(
+      S2AServiceGrpc.S2AServiceStub serviceStub, long deadlineSeconds) {
+    checkNotNull(serviceStub);
+    checkArgument(deadlineSeconds > 0);
+    return new S2AStub(serviceStub, deadlineSeconds);
+  }
+
   @VisibleForTesting
   static S2AStub newInstanceForTesting(StreamObserver<SessionReq> writer) {
     checkNotNull(writer);
@@ -63,6 +72,11 @@ private S2AStub(S2AServiceGrpc.S2AServiceStub serviceStub) {
     this.serviceStub = serviceStub;
   }
 
+  private S2AStub(S2AServiceGrpc.S2AServiceStub serviceStub, long deadlineSeconds) {
+    this.serviceStub = serviceStub;
+    this.deadlineSeconds = deadlineSeconds;
+  }
+
   private S2AStub(StreamObserver<SessionReq> writer) {
     this.writer = writer;
   }
@@ -154,7 +168,7 @@ private void createWriterIfNull() {
       writer =
           serviceStub
               .withWaitForReady()
-              .withDeadlineAfter(HANDSHAKE_RPC_DEADLINE_SECS, SECONDS)
+              .withDeadlineAfter(deadlineSeconds, SECONDS)
               .setUpSession(reader);
     }
   }

s2a/src/test/java/io/grpc/s2a/internal/handshaker/S2AStubTest.java

@@ -61,17 +61,17 @@ public void setUp() {
 
   @Test
   public void send_receiveOkStatus() throws Exception {
-    ObjectPool<Channel> channelPool =
-        SharedResourcePool.forResource(
-            S2AHandshakerServiceChannel.getChannelResource(
-                S2A_ADDRESS, InsecureChannelCredentials.create()));
-    S2AServiceGrpc.S2AServiceStub serviceStub = S2AServiceGrpc.newStub(channelPool.getObject());
-    S2AStub newStub = S2AStub.newInstance(serviceStub);
+    SessionReq req =
+        SessionReq.newBuilder()
+            .setGetTlsConfigurationReq(
+                GetTlsConfigurationReq.newBuilder()
+                    .setConnectionSide(ConnectionSide.CONNECTION_SIDE_CLIENT))
+            .build();
 
-    IOException expected =
-        assertThrows(IOException.class, () -> newStub.send(SessionReq.getDefaultInstance()));
+    SessionResp resp = stub.send(req);
 
-    assertThat(expected).hasMessageThat().contains("DEADLINE_EXCEEDED");
+    assertThat(resp.hasGetTlsConfigurationResp()).isTrue();
+    assertThat(resp.getGetTlsConfigurationResp().hasClientTlsConfiguration()).isTrue();
   }
 
   @Test
@@ -233,6 +233,21 @@ public void send_afterEarlyClose_receivesClosedException() throws InterruptedExc
     assertThat(expected).hasMessageThat().contains("Stream to the S2A is closed.");
   }
 
+  @Test
+  public void send_withUnavailableService_throwsDeadlineExceeded() throws Exception {
+    ObjectPool<Channel> channelPool =
+        SharedResourcePool.forResource(
+            S2AHandshakerServiceChannel.getChannelResource(
+                S2A_ADDRESS, InsecureChannelCredentials.create()));
+    S2AServiceGrpc.S2AServiceStub serviceStub = S2AServiceGrpc.newStub(channelPool.getObject());
+    S2AStub newStub = S2AStub.newInstanceWithDeadline(serviceStub, 1);
+
+    IOException expected =
+        assertThrows(IOException.class, () -> newStub.send(SessionReq.getDefaultInstance()));
+
+    assertThat(expected).hasMessageThat().contains("DEADLINE_EXCEEDED");
+  }
+
   @Test
   public void send_failToWrite() throws Exception {
     FailWriter failWriter = new FailWriter();