Refactor TestRetryingConnectionFactory and TestLazyConnectionFactory

Refactor TestRetryingConnectionFactory and TestLazyConnectionFactory to use Dependency Injection Motivation behind this change is to be able to later add dependent objects to tested class without a need for major changes in the test class itself.
trinodb · Oct 17, 2023 · 4f0f7fb · 4f0f7fb
1 parent aa26d24
commit 4f0f7fb
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 22 deletions.
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/RetryingConnectionFactory.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/RetryingConnectionFactory.java
@@ -14,6 +14,7 @@
 package io.trino.plugin.jdbc;
 
 import com.google.common.base.Throwables;
+import com.google.inject.Inject;
 import dev.failsafe.Failsafe;
 import dev.failsafe.FailsafeException;
 import dev.failsafe.RetryPolicy;
@@ -41,7 +42,8 @@ public class RetryingConnectionFactory
 
     private final ConnectionFactory delegate;
 
-    public RetryingConnectionFactory(ConnectionFactory delegate)
+    @Inject
+    public RetryingConnectionFactory(@StatsCollecting ConnectionFactory delegate)
     {
         this.delegate = requireNonNull(delegate, "delegate is null");
     }

diff --git a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestLazyConnectionFactory.java b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestLazyConnectionFactory.java
@@ -13,6 +13,8 @@
  */
 package io.trino.plugin.jdbc;
 
+import com.google.inject.Guice;
+import com.google.inject.Injector;
 import io.trino.plugin.jdbc.credential.EmptyCredentialProvider;
 import org.h2.Driver;
 import org.junit.jupiter.api.Test;
@@ -30,11 +32,12 @@ public class TestLazyConnectionFactory
     public void testNoConnectionIsCreated()
             throws Exception
     {
-        ConnectionFactory failingConnectionFactory = session -> {
-            throw new AssertionError("Expected no connection creation");
-        };
+        Injector injector = Guice.createInjector(binder -> binder.bind(ConnectionFactory.class).annotatedWith(ForLazyConnectionFactory.class).toInstance(
+                session -> {
+                    throw new AssertionError("Expected no connection creation");
+                }));
 
-        try (LazyConnectionFactory lazyConnectionFactory = new LazyConnectionFactory(failingConnectionFactory);
+        try (LazyConnectionFactory lazyConnectionFactory = injector.getInstance(LazyConnectionFactory.class);
                 Connection ignored = lazyConnectionFactory.openConnection(SESSION)) {
             // no-op
         }
@@ -47,8 +50,10 @@ public void testConnectionCannotBeReusedAfterClose()
         BaseJdbcConfig config = new BaseJdbcConfig()
                 .setConnectionUrl(format("jdbc:h2:mem:test%s;DB_CLOSE_DELAY=-1", System.nanoTime() + ThreadLocalRandom.current().nextLong()));
 
-        try (DriverConnectionFactory h2ConnectionFactory = new DriverConnectionFactory(new Driver(), config, new EmptyCredentialProvider());
-                LazyConnectionFactory lazyConnectionFactory = new LazyConnectionFactory(h2ConnectionFactory)) {
+        Injector injector = Guice.createInjector(binder -> binder.bind(ConnectionFactory.class).annotatedWith(ForLazyConnectionFactory.class).toInstance(
+            new DriverConnectionFactory(new Driver(), config, new EmptyCredentialProvider())));
+
+        try (LazyConnectionFactory lazyConnectionFactory = injector.getInstance(LazyConnectionFactory.class)) {
             Connection connection = lazyConnectionFactory.openConnection(SESSION);
             connection.close();
             assertThatThrownBy(() -> connection.createStatement())

diff --git a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestRetryingConnectionFactory.java b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestRetryingConnectionFactory.java
@@ -13,6 +13,11 @@
  */
 package io.trino.plugin.jdbc;
 
+import com.google.inject.Guice;
+import com.google.inject.Inject;
+import com.google.inject.Injector;
+import com.google.inject.Key;
+import com.google.inject.Scopes;
 import io.trino.spi.StandardErrorCode;
 import io.trino.spi.TrinoException;
 import io.trino.spi.connector.ConnectorSession;
@@ -50,71 +55,102 @@ public void testEverythingImplemented()
     public void testSimplyReturnConnection()
             throws Exception
     {
-        MockConnectorFactory mock = new MockConnectorFactory(RETURN);
-        ConnectionFactory factory = new RetryingConnectionFactory(mock);
-        assertThat(factory.openConnection(SESSION)).isNotNull();
+        Injector injector = createInjector(RETURN);
+        ConnectionFactory factory = injector.getInstance(RetryingConnectionFactory.class);
+        MockConnectorFactory mock = injector.getInstance(MockConnectorFactory.class);
+
+        Connection connection = factory.openConnection(SESSION);
+
+        assertThat(connection).isNotNull();
         assertThat(mock.getCallCount()).isEqualTo(1);
     }
 
     @Test
     public void testRetryAndStopOnTrinoException()
     {
-        MockConnectorFactory mock = new MockConnectorFactory(THROW_SQL_RECOVERABLE_EXCEPTION, THROW_TRINO_EXCEPTION);
-        ConnectionFactory factory = new RetryingConnectionFactory(mock);
+        Injector injector = createInjector(THROW_SQL_RECOVERABLE_EXCEPTION, THROW_TRINO_EXCEPTION);
+        ConnectionFactory factory = injector.getInstance(RetryingConnectionFactory.class);
+        MockConnectorFactory mock = injector.getInstance(MockConnectorFactory.class);
+
         assertThatThrownBy(() -> factory.openConnection(SESSION))
                 .isInstanceOf(TrinoException.class)
                 .hasMessage("Testing Trino exception");
+
         assertThat(mock.getCallCount()).isEqualTo(2);
     }
 
     @Test
     public void testRetryAndStopOnSqlException()
     {
-        MockConnectorFactory mock = new MockConnectorFactory(THROW_SQL_RECOVERABLE_EXCEPTION, THROW_SQL_EXCEPTION);
-        ConnectionFactory factory = new RetryingConnectionFactory(mock);
+        Injector injector = createInjector(THROW_SQL_RECOVERABLE_EXCEPTION, THROW_SQL_EXCEPTION);
+        ConnectionFactory factory = injector.getInstance(RetryingConnectionFactory.class);
+        MockConnectorFactory mock = injector.getInstance(MockConnectorFactory.class);
+
         assertThatThrownBy(() -> factory.openConnection(SESSION))
                 .isInstanceOf(SQLException.class)
                 .hasMessage("Testing sql exception");
+
         assertThat(mock.getCallCount()).isEqualTo(2);
     }
 
     @Test
     public void testNullPointerException()
     {
-        MockConnectorFactory mock = new MockConnectorFactory(THROW_NPE);
-        ConnectionFactory factory = new RetryingConnectionFactory(mock);
+        Injector injector = createInjector(THROW_NPE);
+        ConnectionFactory factory = injector.getInstance(RetryingConnectionFactory.class);
+        MockConnectorFactory mock = injector.getInstance(MockConnectorFactory.class);
+
         assertThatThrownBy(() -> factory.openConnection(SESSION))
                 .isInstanceOf(NullPointerException.class)
                 .hasMessage("Testing NPE");
+
         assertThat(mock.getCallCount()).isEqualTo(1);
     }
 
     @Test
     public void testRetryAndReturn()
             throws Exception
     {
-        MockConnectorFactory mock = new MockConnectorFactory(THROW_SQL_RECOVERABLE_EXCEPTION, RETURN);
-        ConnectionFactory factory = new RetryingConnectionFactory(mock);
-        assertThat(factory.openConnection(SESSION)).isNotNull();
+        Injector injector = createInjector(THROW_SQL_RECOVERABLE_EXCEPTION, RETURN);
+        ConnectionFactory factory = injector.getInstance(RetryingConnectionFactory.class);
+        MockConnectorFactory mock = injector.getInstance(MockConnectorFactory.class);
+
+        Connection connection = factory.openConnection(SESSION);
+
+        assertThat(connection).isNotNull();
         assertThat(mock.getCallCount()).isEqualTo(2);
     }
 
     @Test
     public void testRetryOnWrappedAndReturn()
             throws Exception
     {
-        MockConnectorFactory mock = new MockConnectorFactory(THROW_WRAPPED_SQL_RECOVERABLE_EXCEPTION, RETURN);
-        ConnectionFactory factory = new RetryingConnectionFactory(mock);
-        assertThat(factory.openConnection(SESSION)).isNotNull();
+        Injector injector = createInjector(THROW_WRAPPED_SQL_RECOVERABLE_EXCEPTION, RETURN);
+        ConnectionFactory factory = injector.getInstance(RetryingConnectionFactory.class);
+        MockConnectorFactory mock = injector.getInstance(MockConnectorFactory.class);
+
+        Connection connection = factory.openConnection(SESSION);
+
+        assertThat(connection).isNotNull();
         assertThat(mock.getCallCount()).isEqualTo(2);
     }
 
+    private static Injector createInjector(MockConnectorFactory.Action... actions)
+    {
+        return Guice.createInjector(binder -> {
+            binder.bind(MockConnectorFactory.Action[].class).toInstance(actions);
+            binder.bind(MockConnectorFactory.class).in(Scopes.SINGLETON);
+            binder.bind(ConnectionFactory.class).annotatedWith(StatsCollecting.class).to(Key.get(MockConnectorFactory.class));
+        });
+    }
+
     public static class MockConnectorFactory
             implements ConnectionFactory
     {
         private final Deque<Action> actions = new ArrayDeque<>();
         private int callCount;
 
+        @Inject
         public MockConnectorFactory(Action... actions)
         {
             Stream.of(actions)