add unit test

Author: JoeCqupt

services/src/main/java/io/grpc/protobuf/services/HealthServiceImpl.java

@@ -28,9 +28,11 @@
 import io.grpc.health.v1.HealthCheckResponse.ServingStatus;
 import io.grpc.health.v1.HealthGrpc;
 import io.grpc.stub.StreamObserver;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.IdentityHashMap;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -169,6 +171,18 @@ int numWatchersForTest(String service) {
     }
   }
 
+  @VisibleForTesting
+  Set<StreamObserver<HealthCheckResponse>> watchersForTest(String service) {
+    synchronized (watchLock) {
+      IdentityHashMap<StreamObserver<HealthCheckResponse>, Boolean> serviceWatchers =
+              watchers.get(service);
+      if (serviceWatchers == null) {
+        return Collections.emptySet();
+      }
+      return serviceWatchers.keySet();
+    }
+  }
+
   @GuardedBy("watchLock")
   private void notifyWatchers(String service, @Nullable ServingStatus status) {
     HealthCheckResponse response = getResponseForWatch(status);

services/src/test/java/io/grpc/protobuf/services/HealthStatusManagerTest.java

@@ -31,6 +31,7 @@
 import io.grpc.stub.StreamObserver;
 import io.grpc.testing.GrpcServerRule;
 import java.util.ArrayDeque;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import org.junit.After;
 import org.junit.Before;
@@ -109,6 +110,46 @@ public void enterTerminalState_watch() throws Exception {
     assertThat(obs.responses).isEmpty();
   }
 
+  @Test
+  public void watchNotifyException() throws Exception {
+    manager.setStatus(SERVICE1, ServingStatus.SERVING);
+    CancellableRespObserver cobs1 = new CancellableRespObserver();
+    service.watch(HealthCheckRequest.newBuilder().setService(SERVICE1).build(), cobs1);
+    assertThat(cobs1.responses).hasSize(1);
+    HealthCheckResponse resp = (HealthCheckResponse) cobs1.responses.poll();
+    assertThat(resp.getStatus()).isEqualTo(ServingStatus.SERVING);
+    cobs1.responses.clear();
+
+    CancellableRespObserver cobs2 = new CancellableRespObserver();
+    service.watch(HealthCheckRequest.newBuilder().setService(SERVICE1).build(), cobs2);
+    assertThat(cobs2.responses).hasSize(1);
+    resp = (HealthCheckResponse) cobs2.responses.poll();
+    assertThat(resp.getStatus()).isEqualTo(ServingStatus.SERVING);
+    cobs2.responses.clear();
+
+    // assert that both observers are in watchers
+    Set<StreamObserver<HealthCheckResponse>> observers = service.watchersForTest(SERVICE1);
+    assertThat(observers).hasSize(2);
+
+    // cancel first Observer
+    cobs1 = (CancellableRespObserver) observers.stream().findFirst().get();
+    cobs1.setCancelled(true);
+
+    cobs2 = (CancellableRespObserver) observers.stream().skip(1).findFirst().get();
+
+    manager.setStatus(SERVICE1, ServingStatus.NOT_SERVING);
+
+    // first observer should notify exception
+    assertThat(cobs1.responses).isEmpty();
+
+    // second observer should be notified
+    assertThat(cobs2.responses).hasSize(1);
+    resp = (HealthCheckResponse) cobs2.responses.poll();
+    assertThat(resp.getStatus()).isEqualTo(ServingStatus.NOT_SERVING);
+    cobs2.responses.clear();
+
+  }
+
   @Test
   public void enterTerminalState_ignoreClear() throws Exception {
     manager.setStatus(SERVICE1, ServingStatus.SERVING);
@@ -311,4 +352,23 @@ public void onCompleted() {
       responses.add("onCompleted");
     }
   }
+
+  private static class CancellableRespObserver extends RespObserver {
+
+    boolean cancelled = false;
+
+    public void setCancelled(boolean cancelled) {
+      this.cancelled = cancelled;
+    }
+
+    @Override
+    public void onNext(HealthCheckResponse value) {
+      if (cancelled) {
+        throw Status.CANCELLED
+                .withDescription("call already cancelled.")
+                .asRuntimeException();
+      }
+      super.onNext(value);
+    }
+  }
 }