Fix failing MonitorServiceImplTest tests

Author: aaron-congo

wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java

@@ -42,6 +42,7 @@
 import software.amazon.jdbc.util.ExecutorFactory;
 import software.amazon.jdbc.util.Messages;
 import software.amazon.jdbc.util.PropertyUtils;
+import software.amazon.jdbc.util.connection.ConnectionService;
 import software.amazon.jdbc.util.connection.ConnectionServiceImpl;
 import software.amazon.jdbc.util.events.DataAccessEvent;
 import software.amazon.jdbc.util.events.Event;
@@ -197,20 +198,15 @@ public <T extends Monitor> T runIfAbsent(
       cacheContainer = monitorCaches.computeIfAbsent(monitorClass, k -> supplier.get());
     }
 
-    TargetDriverHelper helper = new TargetDriverHelper();
-    java.sql.Driver driver = helper.getTargetDriver(originalUrl, originalProps);
-    final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver);
-    final Properties propsCopy = PropertyUtils.copyProperties(originalProps);
-    final ConnectionServiceImpl connectionService = new ConnectionServiceImpl(
-        storageService,
-        this,
-        telemetryFactory,
-        defaultConnectionProvider,
-        originalUrl,
-        driverProtocol,
-        driverDialect,
-        dbDialect,
-        propsCopy);
+    final ConnectionService connectionService =
+        getConnectionService(
+            storageService,
+            telemetryFactory,
+            originalUrl,
+            driverProtocol,
+            driverDialect,
+            dbDialect,
+            originalProps);
 
     Monitor monitor = cacheContainer.getCache().computeIfAbsent(key, k -> {
       MonitorItem monitorItem = new MonitorItem(() -> initializer.createMonitor(
@@ -228,6 +224,25 @@ public <T extends Monitor> T runIfAbsent(
         Messages.get("MonitorServiceImpl.unexpectedMonitorClass", new Object[] {monitorClass, monitor}));
   }
 
+  protected ConnectionService getConnectionService(StorageService storageService,
+      TelemetryFactory telemetryFactory, String originalUrl, String driverProtocol, TargetDriverDialect driverDialect,
+      Dialect dbDialect, Properties originalProps) throws SQLException {
+    TargetDriverHelper helper = new TargetDriverHelper();
+    java.sql.Driver driver = helper.getTargetDriver(originalUrl, originalProps);
+    final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver);
+    final Properties propsCopy = PropertyUtils.copyProperties(originalProps);
+    return new ConnectionServiceImpl(
+        storageService,
+        this,
+        telemetryFactory,
+        defaultConnectionProvider,
+        originalUrl,
+        driverProtocol,
+        driverDialect,
+        dbDialect,
+        propsCopy);
+  }
+
   @Override
   public @Nullable <T extends Monitor> T get(Class<T> monitorClass, Object key) {
     CacheContainer cacheContainer = monitorCaches.get(monitorClass);

wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java

@@ -21,86 +21,98 @@
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
 
 import java.sql.SQLException;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Properties;
 import java.util.concurrent.TimeUnit;
 import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import software.amazon.jdbc.dialect.Dialect;
 import software.amazon.jdbc.plugin.customendpoint.CustomEndpointMonitorImpl;
 import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect;
+import software.amazon.jdbc.util.connection.ConnectionService;
 import software.amazon.jdbc.util.events.EventPublisher;
 import software.amazon.jdbc.util.storage.StorageService;
 import software.amazon.jdbc.util.telemetry.TelemetryFactory;
 
 class MonitorServiceImplTest {
-  @Mock StorageService storageService;
-  @Mock TelemetryFactory telemetryFactory;
-  @Mock TargetDriverDialect targetDriverDialect;
-  @Mock Dialect dbDialect;
-  @Mock EventPublisher publisher;
-  MonitorServiceImpl monitorService;
+  @Mock StorageService mockStorageService;
+  @Mock ConnectionService mockConnectionService;
+  @Mock TelemetryFactory mockTelemetryFactory;
+  @Mock TargetDriverDialect mockTargetDriverDialect;
+  @Mock Dialect mockDbDialect;
+  @Mock EventPublisher mockPublisher;
+  MonitorServiceImpl spyMonitorService;
   private AutoCloseable closeable;
 
   @BeforeEach
   void setUp() {
     closeable = MockitoAnnotations.openMocks(this);
-    monitorService = new MonitorServiceImpl(publisher) {
-      @Override
-      protected void initCleanupThread(long cleanupIntervalNanos) {
-        // Do nothing
-      }
-    };
+    spyMonitorService = spy(new MonitorServiceImpl(mockPublisher));
+    doNothing().when(spyMonitorService).initCleanupThread(anyInt());
+
+    try {
+      doReturn(mockConnectionService).when(spyMonitorService)
+          .getConnectionService(any(), any(), any(), any(), any(), any(), any());
+    } catch (SQLException e) {
+      Assertions.fail(
+          "Encountered exception while stubbing MonitorServiceImpl#getConnectionService: " + e.getMessage());
+    }
   }
 
   @AfterEach
   void tearDown() throws Exception {
     closeable.close();
-    monitorService.releaseResources();
+    spyMonitorService.releaseResources();
   }
 
   @Test
   public void testMonitorError_monitorReCreated() throws SQLException, InterruptedException {
-    monitorService.registerMonitorTypeIfAbsent(
+    spyMonitorService.registerMonitorTypeIfAbsent(
         NoOpMonitor.class,
         TimeUnit.MINUTES.toNanos(1),
         TimeUnit.MINUTES.toNanos(1),
         new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)),
         null
     );
     String key = "testMonitor";
-    NoOpMonitor monitor = monitorService.runIfAbsent(
+    NoOpMonitor monitor = spyMonitorService.runIfAbsent(
         NoOpMonitor.class,
         key,
-        storageService,
-        telemetryFactory,
+        mockStorageService,
+        mockTelemetryFactory,
         "jdbc:postgresql://somehost/somedb",
         "someProtocol",
-        targetDriverDialect,
-        dbDialect,
+        mockTargetDriverDialect,
+        mockDbDialect,
         new Properties(),
-        (connectionService, pluginService) -> new NoOpMonitor(monitorService, 30)
+        (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30)
     );
 
-    Monitor storedMonitor = monitorService.get(NoOpMonitor.class, key);
+    Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key);
     assertNotNull(storedMonitor);
     assertEquals(monitor, storedMonitor);
     // need to wait to give time for the monitor executor to start the monitor thread.
     TimeUnit.MILLISECONDS.sleep(250);
     assertEquals(MonitorState.RUNNING, monitor.getState());
 
     monitor.state.set(MonitorState.ERROR);
-    monitorService.checkMonitors();
+    spyMonitorService.checkMonitors();
 
     assertEquals(MonitorState.STOPPED, monitor.getState());
 
-    Monitor newMonitor = monitorService.get(NoOpMonitor.class, key);
+    Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key);
     assertNotNull(newMonitor);
     assertNotEquals(monitor, newMonitor);
     // need to wait to give time for the monitor executor to start the monitor thread.
@@ -110,40 +122,40 @@ public void testMonitorError_monitorReCreated() throws SQLException, Interrupted
 
   @Test
   public void testMonitorStuck_monitorReCreated() throws SQLException, InterruptedException {
-    monitorService.registerMonitorTypeIfAbsent(
+    spyMonitorService.registerMonitorTypeIfAbsent(
         NoOpMonitor.class,
         TimeUnit.MINUTES.toNanos(1),
         1, // heartbeat times out immediately
         new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)),
         null
     );
     String key = "testMonitor";
-    NoOpMonitor monitor = monitorService.runIfAbsent(
+    NoOpMonitor monitor = spyMonitorService.runIfAbsent(
         NoOpMonitor.class,
         key,
-        storageService,
-        telemetryFactory,
+        mockStorageService,
+        mockTelemetryFactory,
         "jdbc:postgresql://somehost/somedb",
         "someProtocol",
-        targetDriverDialect,
-        dbDialect,
+        mockTargetDriverDialect,
+        mockDbDialect,
         new Properties(),
-        (connectionService, pluginService) -> new NoOpMonitor(monitorService, 30)
+        (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30)
     );
 
-    Monitor storedMonitor = monitorService.get(NoOpMonitor.class, key);
+    Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key);
     assertNotNull(storedMonitor);
     assertEquals(monitor, storedMonitor);
     // need to wait to give time for the monitor executor to start the monitor thread.
     TimeUnit.MILLISECONDS.sleep(250);
     assertEquals(MonitorState.RUNNING, monitor.getState());
 
     // checkMonitors() should detect the heartbeat/inactivity timeout, stop the monitor, and re-create a new one.
-    monitorService.checkMonitors();
+    spyMonitorService.checkMonitors();
 
     assertEquals(MonitorState.STOPPED, monitor.getState());
 
-    Monitor newMonitor = monitorService.get(NoOpMonitor.class, key);
+    Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key);
     assertNotNull(newMonitor);
     assertNotEquals(monitor, newMonitor);
     // need to wait to give time for the monitor executor to start the monitor thread.
@@ -153,7 +165,7 @@ public void testMonitorStuck_monitorReCreated() throws SQLException, Interrupted
 
   @Test
   public void testMonitorExpired() throws SQLException, InterruptedException {
-    monitorService.registerMonitorTypeIfAbsent(
+    spyMonitorService.registerMonitorTypeIfAbsent(
         NoOpMonitor.class,
         TimeUnit.MILLISECONDS.toNanos(200), // monitor expires after 200ms
         TimeUnit.MINUTES.toNanos(1),
@@ -163,57 +175,57 @@ public void testMonitorExpired() throws SQLException, InterruptedException {
         null
     );
     String key = "testMonitor";
-    NoOpMonitor monitor = monitorService.runIfAbsent(
+    NoOpMonitor monitor = spyMonitorService.runIfAbsent(
         NoOpMonitor.class,
         key,
-        storageService,
-        telemetryFactory,
+        mockStorageService,
+        mockTelemetryFactory,
         "jdbc:postgresql://somehost/somedb",
         "someProtocol",
-        targetDriverDialect,
-        dbDialect,
+        mockTargetDriverDialect,
+        mockDbDialect,
         new Properties(),
-        (connectionService, pluginService) -> new NoOpMonitor(monitorService, 30)
+        (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30)
     );
 
-    Monitor storedMonitor = monitorService.get(NoOpMonitor.class, key);
+    Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key);
     assertNotNull(storedMonitor);
     assertEquals(monitor, storedMonitor);
     // need to wait to give time for the monitor executor to start the monitor thread.
     TimeUnit.MILLISECONDS.sleep(250);
     assertEquals(MonitorState.RUNNING, monitor.getState());
 
     // checkMonitors() should detect the expiration timeout and stop/remove the monitor.
-    monitorService.checkMonitors();
+    spyMonitorService.checkMonitors();
 
     assertEquals(MonitorState.STOPPED, monitor.getState());
 
-    Monitor newMonitor = monitorService.get(NoOpMonitor.class, key);
+    Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key);
     // monitor should have been removed when checkMonitors() was called.
     assertNull(newMonitor);
   }
 
   @Test
   public void testMonitorMismatch() {
-    assertThrows(IllegalStateException.class, () -> monitorService.runIfAbsent(
+    assertThrows(IllegalStateException.class, () -> spyMonitorService.runIfAbsent(
         CustomEndpointMonitorImpl.class,
         "testMonitor",
-        storageService,
-        telemetryFactory,
+        mockStorageService,
+        mockTelemetryFactory,
         "jdbc:postgresql://somehost/somedb",
         "someProtocol",
-        targetDriverDialect,
-        dbDialect,
+        mockTargetDriverDialect,
+        mockDbDialect,
         new Properties(),
         // indicated monitor class is CustomEndpointMonitorImpl, but actual monitor is NoOpMonitor. The monitor
         // service should detect this and throw an exception.
-        (connectionService, pluginService) -> new NoOpMonitor(monitorService, 30)
+        (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30)
     ));
   }
 
   @Test
   public void testRemove() throws SQLException, InterruptedException {
-    monitorService.registerMonitorTypeIfAbsent(
+    spyMonitorService.registerMonitorTypeIfAbsent(
         NoOpMonitor.class,
         TimeUnit.MINUTES.toNanos(1),
         TimeUnit.MINUTES.toNanos(1),
@@ -224,30 +236,30 @@ public void testRemove() throws SQLException, InterruptedException {
     );
 
     String key = "testMonitor";
-    NoOpMonitor monitor = monitorService.runIfAbsent(
+    NoOpMonitor monitor = spyMonitorService.runIfAbsent(
         NoOpMonitor.class,
         key,
-        storageService,
-        telemetryFactory,
+        mockStorageService,
+        mockTelemetryFactory,
         "jdbc:postgresql://somehost/somedb",
         "someProtocol",
-        targetDriverDialect,
-        dbDialect,
+        mockTargetDriverDialect,
+        mockDbDialect,
         new Properties(),
-        (connectionService, pluginService) -> new NoOpMonitor(monitorService, 30)
+        (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30)
     );
     assertNotNull(monitor);
 
     // need to wait to give time for the monitor executor to start the monitor thread.
     TimeUnit.MILLISECONDS.sleep(250);
-    Monitor removedMonitor = monitorService.remove(NoOpMonitor.class, key);
+    Monitor removedMonitor = spyMonitorService.remove(NoOpMonitor.class, key);
     assertEquals(monitor, removedMonitor);
     assertEquals(MonitorState.RUNNING, monitor.getState());
   }
 
   @Test
   public void testStopAndRemove() throws SQLException, InterruptedException {
-    monitorService.registerMonitorTypeIfAbsent(
+    spyMonitorService.registerMonitorTypeIfAbsent(
         NoOpMonitor.class,
         TimeUnit.MINUTES.toNanos(1),
         TimeUnit.MINUTES.toNanos(1),
@@ -258,24 +270,24 @@ public void testStopAndRemove() throws SQLException, InterruptedException {
     );
 
     String key = "testMonitor";
-    NoOpMonitor monitor = monitorService.runIfAbsent(
+    NoOpMonitor monitor = spyMonitorService.runIfAbsent(
         NoOpMonitor.class,
         key,
-        storageService,
-        telemetryFactory,
+        mockStorageService,
+        mockTelemetryFactory,
         "jdbc:postgresql://somehost/somedb",
         "someProtocol",
-        targetDriverDialect,
-        dbDialect,
+        mockTargetDriverDialect,
+        mockDbDialect,
         new Properties(),
-        (connectionService, pluginService) -> new NoOpMonitor(monitorService, 30)
+        (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30)
     );
     assertNotNull(monitor);
 
     // need to wait to give time for the monitor executor to start the monitor thread.
     TimeUnit.MILLISECONDS.sleep(250);
-    monitorService.stopAndRemove(NoOpMonitor.class, key);
-    assertNull(monitorService.get(NoOpMonitor.class, key));
+    spyMonitorService.stopAndRemove(NoOpMonitor.class, key);
+    assertNull(spyMonitorService.get(NoOpMonitor.class, key));
     assertEquals(MonitorState.STOPPED, monitor.getState());
   }