From fb4a18034ac8b5eac0fb61e1b2748af7748f7dbb Mon Sep 17 00:00:00 2001
From: Blake Li <blakeli@google.com>
Date: Fri, 2 Dec 2022 21:58:07 +0000
Subject: [PATCH] test: Add unit tests for Watchdog (#1919)

---
 .../com/google/api/gax/rpc/WatchdogTest.java  | 81 +++++++++++++++++--
 1 file changed, 73 insertions(+), 8 deletions(-)

diff --git a/gax/src/test/java/com/google/api/gax/rpc/WatchdogTest.java b/gax/src/test/java/com/google/api/gax/rpc/WatchdogTest.java
index 0b161018f..e20218452 100644
--- a/gax/src/test/java/com/google/api/gax/rpc/WatchdogTest.java
+++ b/gax/src/test/java/com/google/api/gax/rpc/WatchdogTest.java
@@ -44,6 +44,7 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -195,14 +196,7 @@ public void testMultiple() throws Exception {
   @SuppressWarnings("unchecked")
   public void testWatchdogBeingClosed() {
     ScheduledFuture future = Mockito.mock(ScheduledFuture.class);
-    ScheduledExecutorService mockExecutor = Mockito.mock(ScheduledExecutorService.class);
-    Mockito.when(
-            mockExecutor.scheduleAtFixedRate(
-                Mockito.any(Watchdog.class),
-                Mockito.anyLong(),
-                Mockito.anyLong(),
-                Mockito.any(TimeUnit.class)))
-        .thenReturn(future);
+    ScheduledExecutorService mockExecutor = getMockExecutorService(future);
     Watchdog underTest = Watchdog.create(clock, checkInterval, mockExecutor);
     assertThat(underTest).isInstanceOf(BackgroundResource.class);
 
@@ -219,6 +213,77 @@ public void testWatchdogBeingClosed() {
     Mockito.verifyNoMoreInteractions(mockExecutor);
   }
 
+  @Test
+  public void awaitTermination_shouldReturnTrueIfFutureIsDone() throws Exception {
+    int duration = 1000;
+    TimeUnit timeUnit = TimeUnit.MILLISECONDS;
+    ScheduledFuture future = Mockito.mock(ScheduledFuture.class);
+    ScheduledExecutorService mockExecutor = getMockExecutorService(future);
+    Watchdog watchdog = Watchdog.create(clock, checkInterval, mockExecutor);
+    watchdog.shutdown();
+
+    boolean actual = watchdog.awaitTermination(duration, timeUnit);
+
+    assertThat(actual).isTrue();
+  }
+
+  @Test
+  public void awaitTermination_shouldReturnFalseIfGettingFutureTimedOut() throws Exception {
+    int duration = 1000;
+    TimeUnit timeUnit = TimeUnit.MILLISECONDS;
+    ScheduledFuture future = Mockito.mock(ScheduledFuture.class);
+    Mockito.doThrow(new TimeoutException()).when(future).get(duration, timeUnit);
+    ScheduledExecutorService mockExecutor = getMockExecutorService(future);
+    Watchdog watchdog = Watchdog.create(clock, checkInterval, mockExecutor);
+
+    boolean actual = watchdog.awaitTermination(duration, timeUnit);
+
+    assertThat(actual).isFalse();
+  }
+
+  @Test
+  public void awaitTermination_shouldReturnTrueIfFutureIsAlreadyCancelled() throws Exception {
+    int duration = 1000;
+    TimeUnit timeUnit = TimeUnit.MILLISECONDS;
+    ScheduledFuture future = Mockito.mock(ScheduledFuture.class);
+    Mockito.doThrow(new CancellationException()).when(future).get(duration, timeUnit);
+    ScheduledExecutorService mockExecutor = getMockExecutorService(future);
+    Watchdog watchdog = Watchdog.create(clock, checkInterval, mockExecutor);
+
+    boolean actual = watchdog.awaitTermination(duration, timeUnit);
+
+    assertThat(actual).isTrue();
+  }
+
+  @Test
+  public void awaitTermination_shouldReturnFalseIfGettingFutureThrowsExecutionException()
+      throws Exception {
+    int duration = 1000;
+    TimeUnit timeUnit = TimeUnit.MILLISECONDS;
+    ScheduledFuture future = Mockito.mock(ScheduledFuture.class);
+    Mockito.doThrow(new ExecutionException(new RuntimeException()))
+        .when(future)
+        .get(duration, timeUnit);
+    ScheduledExecutorService mockExecutor = getMockExecutorService(future);
+    Watchdog watchdog = Watchdog.create(clock, checkInterval, mockExecutor);
+
+    boolean actual = watchdog.awaitTermination(duration, timeUnit);
+
+    assertThat(actual).isTrue();
+  }
+
+  private ScheduledExecutorService getMockExecutorService(ScheduledFuture future) {
+    ScheduledExecutorService mockExecutor = Mockito.mock(ScheduledExecutorService.class);
+    Mockito.when(
+            mockExecutor.scheduleAtFixedRate(
+                Mockito.any(Watchdog.class),
+                Mockito.anyLong(),
+                Mockito.anyLong(),
+                Mockito.any(TimeUnit.class)))
+        .thenReturn(future);
+    return mockExecutor;
+  }
+
   static class AccumulatingObserver<T> implements ResponseObserver<T> {
     SettableApiFuture<StreamController> controller = SettableApiFuture.create();
     Queue<T> responses = Queues.newLinkedBlockingDeque();