test: clean up CookieHolderTest (#2361)

* test: clean up CookieHolderTest

- stop returning protos to trigger headers
- fix thread safety for variables that are accessed by multiple threads

Change-Id: Ib34440d2ffeddcb5611b5e1cee183004bafc9449

* missing import

Change-Id: Ied61bc081c3cc0220fd96051aabd1849cbffd10f

* one more

Change-Id: Ibf2f4ad8bea038390b1ddf099c3061f0d9b02495

* fix broken test

Change-Id: I63a6e54829ef5ea2db8d3b3b2db0afdc89cc1248

* make it consistent with other tests

Change-Id: I2296b1cb7f52895a59e6b4fe5586d07c0dfd97a2

Author: igorbernstein2

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java

@@ -63,6 +63,7 @@
 import io.grpc.stub.StreamObserver;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
@@ -94,9 +95,9 @@ public class CookiesHolderTest {
   private final FakeService fakeService = new FakeService();
   private BigtableDataSettings.Builder settings;
   private BigtableDataClient client;
-  private final List<Metadata> serverMetadata = new ArrayList<>();
+  private final List<Metadata> serverMetadata = Collections.synchronizedList(new ArrayList<>());
 
-  private final Set<String> methods = new HashSet<>();
+  private final Set<String> methods = Collections.synchronizedSet(new HashSet<>());
 
   @Before
   public void setup() throws Exception {
@@ -111,13 +112,17 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
             if (metadata.containsKey(ROUTING_COOKIE_1)) {
               methods.add(serverCall.getMethodDescriptor().getBareMethodName());
             }
+
+            Metadata responseHeaders = new Metadata();
+            responseHeaders.put(ROUTING_COOKIE_HEADER, testHeaderCookie);
+            responseHeaders.put(ROUTING_COOKIE_1, routingCookie1Header);
+            serverCall.sendHeaders(responseHeaders);
+
             return serverCallHandler.startCall(
                 new ForwardingServerCall.SimpleForwardingServerCall<ReqT, RespT>(serverCall) {
                   @Override
                   public void sendHeaders(Metadata responseHeaders) {
-                    responseHeaders.put(ROUTING_COOKIE_HEADER, testHeaderCookie);
-                    responseHeaders.put(ROUTING_COOKIE_1, routingCookie1Header);
-                    super.sendHeaders(responseHeaders);
+                    // headers already sent!
                   }
                 },
                 metadata);
@@ -450,8 +455,8 @@ public void testNoCookieSucceedGenerateInitialChangeStreamParition() {
 
     Metadata lastMetadata = serverMetadata.get(fakeService.count.get() - 1);
 
-    assertThat(lastMetadata)
-        .doesNotContainKeys(ROUTING_COOKIE_1.name(), ROUTING_COOKIE_2.name(), BAD_KEY.name());
+    assertThat(lastMetadata).doesNotContainKeys(ROUTING_COOKIE_2.name(), BAD_KEY.name());
+    assertThat(lastMetadata).containsAtLeast(ROUTING_COOKIE_1.name(), routingCookie1Header);
 
     serverMetadata.clear();
   }
@@ -657,7 +662,7 @@ public void testDisableRoutingCookie() throws IOException {
 
   static class FakeService extends BigtableGrpc.BigtableImplBase {
 
-    private boolean returnCookie = true;
+    private volatile boolean returnCookie = true;
     private final AtomicInteger count = new AtomicInteger();
 
     @Override
@@ -666,7 +671,6 @@ public void readRows(
       if (count.getAndIncrement() < 1) {
         Metadata trailers = new Metadata();
         maybePopulateCookie(trailers, "readRows");
-        responseObserver.onNext(ReadRowsResponse.getDefaultInstance());
         StatusRuntimeException exception = new StatusRuntimeException(Status.UNAVAILABLE, trailers);
         responseObserver.onError(exception);
         return;
@@ -681,7 +685,6 @@ public void mutateRow(
       if (count.getAndIncrement() < 1) {
         Metadata trailers = new Metadata();
         maybePopulateCookie(trailers, "mutateRow");
-        responseObserver.onNext(MutateRowResponse.getDefaultInstance());
         StatusRuntimeException exception = new StatusRuntimeException(Status.UNAVAILABLE, trailers);
         responseObserver.onError(exception);
         return;
@@ -696,7 +699,6 @@ public void mutateRows(
       if (count.getAndIncrement() < 1) {
         Metadata trailers = new Metadata();
         maybePopulateCookie(trailers, "mutateRows");
-        responseObserver.onNext(MutateRowsResponse.getDefaultInstance());
         StatusRuntimeException exception = new StatusRuntimeException(Status.UNAVAILABLE, trailers);
         responseObserver.onError(exception);
         return;
@@ -714,7 +716,6 @@ public void sampleRowKeys(
       if (count.getAndIncrement() < 1) {
         Metadata trailers = new Metadata();
         maybePopulateCookie(trailers, "sampleRowKeys");
-        responseObserver.onNext(SampleRowKeysResponse.getDefaultInstance());
         StatusRuntimeException exception = new StatusRuntimeException(Status.UNAVAILABLE, trailers);
         responseObserver.onError(exception);
         return;