Improve multi-server connection testing (#3170)

ohadzeliger · web-flow · commit 803f3d74e3d0 · 2025-02-20T16:32:17.000-05:00
Add tests that assert on the returned statements from the multi-server-connection Change connection logic to keep same current selector on read operations Resolves #3094
diff --git a/yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/MultiServerConnectionFactory.java b/yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/MultiServerConnectionFactory.java
@@ -174,12 +174,12 @@ public int getCurrentConnectionSelector() {
 
         @Override
         public RelationalStatement createStatement() throws SQLException {
-            return getNextConnection("createStatement").createStatement();
+            return getCurrentConnection(true, "createStatement").createStatement();
         }
 
         @Override
         public RelationalPreparedStatement prepareStatement(String sql) throws SQLException {
-            return getNextConnection("prepareStatement").prepareStatement(sql);
+            return getCurrentConnection(true, "prepareStatement").prepareStatement(sql);
         }
 
         @Override
@@ -195,7 +195,7 @@ public void setAutoCommit(boolean autoCommit) throws SQLException {
 
         @Override
         public boolean getAutoCommit() throws SQLException {
-            return getNextConnection("getAutoCommit").getAutoCommit();
+            return getCurrentConnection(false, "getAutoCommit").getAutoCommit();
         }
 
         @Override
@@ -218,32 +218,35 @@ public void close() throws SQLException {
 
         @Override
         public boolean isClosed() throws SQLException {
-            return getNextConnection("isClosed").isClosed();
+            return getCurrentConnection(false, "isClosed").isClosed();
         }
 
         @Override
         public DatabaseMetaData getMetaData() throws SQLException {
-            return getNextConnection("getMetaData").getMetaData();
+            return getCurrentConnection(false, "getMetaData").getMetaData();
         }
 
         @Override
         public void setTransactionIsolation(int level) throws SQLException {
-            getNextConnection("setTransactionIsolation").setTransactionIsolation(level);
+            logger.info("Sending operation {} to all connections", "setTransactionIsolation");
+            for (RelationalConnection connection : relationalConnections) {
+                connection.setTransactionIsolation(level);
+            }
         }
 
         @Override
         public int getTransactionIsolation() throws SQLException {
-            return getNextConnection("getTransactionIsolation").getTransactionIsolation();
+            return getCurrentConnection(false, "getTransactionIsolation").getTransactionIsolation();
         }
 
         @Override
         public Array createArrayOf(String typeName, Object[] elements) throws SQLException {
-            return getNextConnection("createArrayOf").createArrayOf(typeName, elements);
+            return getCurrentConnection(true, "createArrayOf").createArrayOf(typeName, elements);
         }
 
         @Override
         public Struct createStruct(String typeName, Object[] attributes) throws SQLException {
-            return getNextConnection("createStruct").createStruct(typeName, attributes);
+            return getCurrentConnection(true, "createStruct").createStruct(typeName, attributes);
         }
 
         @Override
@@ -256,13 +259,13 @@ public void setSchema(String schema) throws SQLException {
 
         @Override
         public String getSchema() throws SQLException {
-            return getNextConnection("getSchema").getSchema();
+            return getCurrentConnection(false, "getSchema").getSchema();
         }
 
         @Nonnull
         @Override
         public Options getOptions() {
-            return getNextConnection("getOptions").getOptions();
+            return getCurrentConnection(false, "getOptions").getOptions();
         }
 
         @Override
@@ -275,10 +278,19 @@ public void setOption(Options.Name name, Object value) throws SQLException {
 
         @Override
         public URI getPath() {
-            return getNextConnection("getPath").getPath();
+            return getCurrentConnection(false, "getPath").getPath();
         }
 
-        private RelationalConnection getNextConnection(String op) {
+        /**
+         * Get the connection to send requests to.
+         * This method conditionally advances the connection selector. This is done for ease of testing, where some
+         * commands that just return information will not advance the selector, leading to better predictability
+         * of the command routing.
+         * @param advance whether to advance the connection selector
+         * @param op the name of the operation (for logging)
+         * @return the underlying connection to use
+         */
+        private RelationalConnection getCurrentConnection(boolean advance, String op) {
             switch (connectionSelectionPolicy) {
                 case DEFAULT:
                     if (logger.isInfoEnabled()) {
@@ -290,7 +302,9 @@ private RelationalConnection getNextConnection(String op) {
                     if (logger.isInfoEnabled()) {
                         logger.info("Sending operation {} to connection {}", op, currentConnectionSelector);
                     }
-                    currentConnectionSelector = (currentConnectionSelector + 1) % relationalConnections.size();
+                    if (advance) {
+                        currentConnectionSelector = (currentConnectionSelector + 1) % relationalConnections.size();
+                    }
                     return result;
                 default:
                     throw new IllegalStateException("Unsupported selection policy " + connectionSelectionPolicy);
diff --git a/yaml-tests/src/test/java/MultiServerConnectionFactoryTest.java b/yaml-tests/src/test/java/MultiServerConnectionFactoryTest.java
@@ -19,6 +19,7 @@
  */
 
 import com.apple.foundationdb.relational.api.RelationalConnection;
+import com.apple.foundationdb.relational.api.RelationalPreparedStatement;
 import com.apple.foundationdb.relational.yamltests.MultiServerConnectionFactory;
 import com.apple.foundationdb.relational.yamltests.YamlRunner;
 import org.junit.jupiter.api.Test;
@@ -42,61 +43,110 @@ void testDefaultPolicy(int initialConnection) throws SQLException {
         MultiServerConnectionFactory classUnderTest = new MultiServerConnectionFactory(
                 MultiServerConnectionFactory.ConnectionSelectionPolicy.DEFAULT,
                 initialConnection,
-                dummyConnectionFactory(),
-                List.of(dummyConnectionFactory()));
-        assertEquals(initialConnection, classUnderTest.getCurrentConnectionSelector());
-        MultiServerConnectionFactory.MultiServerRelationalConnection connection = (MultiServerConnectionFactory.MultiServerRelationalConnection)classUnderTest.getNewConnection(URI.create("Blah"));
-        assertEquals(0, connection.getCurrentConnectionSelector());
-        assertEquals(initialConnection, classUnderTest.getCurrentConnectionSelector());
+                dummyConnectionFactory("Primary"),
+                List.of(dummyConnectionFactory("Alternate")));
+
+        MultiServerConnectionFactory.MultiServerRelationalConnection connection =
+                (MultiServerConnectionFactory.MultiServerRelationalConnection)classUnderTest.getNewConnection(URI.create("Blah"));
+        assertConnection(connection, "Primary");
+        assertStatement(connection.prepareStatement("SQL"), "Primary");
+        assertStatement(connection.prepareStatement("SQL"), "Primary");
+
         connection = (MultiServerConnectionFactory.MultiServerRelationalConnection)classUnderTest.getNewConnection(URI.create("Blah"));
-        assertEquals(0, connection.getCurrentConnectionSelector());
-        assertEquals(initialConnection, classUnderTest.getCurrentConnectionSelector());
+        assertConnection(connection, "Primary");
+        assertStatement(connection.prepareStatement("SQL"), "Primary");
+        assertStatement(connection.prepareStatement("SQL"), "Primary");
+
         connection = (MultiServerConnectionFactory.MultiServerRelationalConnection)classUnderTest.getNewConnection(URI.create("Blah"));
-        assertEquals(0, connection.getCurrentConnectionSelector());
-        assertEquals(initialConnection, classUnderTest.getCurrentConnectionSelector());
+        assertConnection(connection, "Primary");
+        assertStatement(connection.prepareStatement("SQL"), "Primary");
+        assertStatement(connection.prepareStatement("SQL"), "Primary");
     }
 
     @ParameterizedTest
     @CsvSource({"0", "1"})
     void testAlternatePolicy(int initialConnection) throws SQLException {
+        final String[] path = new String[] { "Primary", "Alternate" };
+        final String initialConnectionName = path[initialConnection];
+        final String otherConnectionName = path[(initialConnection + 1) % 2];
+
         MultiServerConnectionFactory classUnderTest = new MultiServerConnectionFactory(
                 MultiServerConnectionFactory.ConnectionSelectionPolicy.ALTERNATE,
                 initialConnection,
-                dummyConnectionFactory(),
-                List.of(dummyConnectionFactory()));
-        assertEquals(initialConnection, classUnderTest.getCurrentConnectionSelector());
-        MultiServerConnectionFactory.MultiServerRelationalConnection connection = (MultiServerConnectionFactory.MultiServerRelationalConnection)classUnderTest.getNewConnection(URI.create("Blah"));
-        assertEquals(initialConnection % 2, connection.getCurrentConnectionSelector());
-        assertEquals((initialConnection + 1) % 2, classUnderTest.getCurrentConnectionSelector());
+                dummyConnectionFactory("Primary"),
+                List.of(dummyConnectionFactory("Alternate")));
+
+        // First run:
+        // - Factory current connection: initial connection
+        // - connection current connection: initial connection
+        // - statement: initial connection (2 statements)
+        MultiServerConnectionFactory.MultiServerRelationalConnection connection =
+                (MultiServerConnectionFactory.MultiServerRelationalConnection)classUnderTest.getNewConnection(URI.create("Blah"));
+        assertConnection(connection, initialConnectionName);
+        assertStatement(connection.prepareStatement("SQL"), initialConnectionName);
+        // next statement
+        assertStatement(connection.prepareStatement("SQL"), otherConnectionName);
+
+        // Second run:
+        // - Factory current connection: alternate connection
+        // - connection current connection: alternate connection
+        // - statement: alternate connection (2 statements)
+        connection = (MultiServerConnectionFactory.MultiServerRelationalConnection)classUnderTest.getNewConnection(URI.create("Blah"));
+        assertConnection(connection, otherConnectionName);
+        assertStatement(connection.prepareStatement("SQL"), otherConnectionName);
+        // next statement
+        assertStatement(connection.prepareStatement("SQL"), initialConnectionName);
+
+        // Third run:
+        // - Factory current connection: initial connection
+        // - connection current connection: initial connection
+        // - statement: initial connection (1 statement)
         connection = (MultiServerConnectionFactory.MultiServerRelationalConnection)classUnderTest.getNewConnection(URI.create("Blah"));
-        assertEquals((initialConnection + 1) % 2, connection.getCurrentConnectionSelector());
-        assertEquals(initialConnection % 2, classUnderTest.getCurrentConnectionSelector());
+        assertConnection(connection, initialConnectionName);
+        // just one statement for this connection
+        assertStatement(connection.prepareStatement("SQL"), initialConnectionName);
+
+        // Fourth run:
+        // - Factory current connection: alternate connection
+        // - connection current connection: alternate connection
+        // - statement: alternate connection (3 statements)
         connection = (MultiServerConnectionFactory.MultiServerRelationalConnection)classUnderTest.getNewConnection(URI.create("Blah"));
-        assertEquals(initialConnection % 2, connection.getCurrentConnectionSelector());
-        assertEquals((initialConnection + 1) % 2, classUnderTest.getCurrentConnectionSelector());
+        assertConnection(connection, otherConnectionName);
+        assertStatement(connection.prepareStatement("SQL"), otherConnectionName);
+        // next statements
+        assertStatement(connection.prepareStatement("SQL"), initialConnectionName);
+        assertStatement(connection.prepareStatement("SQL"), otherConnectionName);
     }
 
     @Test
     void testIllegalInitialConnection() {
         assertThrows(AssertionError.class, () -> new MultiServerConnectionFactory(
                 MultiServerConnectionFactory.ConnectionSelectionPolicy.ALTERNATE,
                 -1,
-                dummyConnectionFactory(),
-                List.of(dummyConnectionFactory())));
+                dummyConnectionFactory("A"),
+                List.of(dummyConnectionFactory("B"))));
         assertThrows(AssertionError.class, () -> new MultiServerConnectionFactory(
                 MultiServerConnectionFactory.ConnectionSelectionPolicy.ALTERNATE,
                 7,
-                dummyConnectionFactory(),
-                List.of(dummyConnectionFactory())));
+                dummyConnectionFactory("A"),
+                List.of(dummyConnectionFactory("B"))));
     }
 
-    YamlRunner.YamlConnectionFactory dummyConnectionFactory() {
+    private void assertStatement(final RelationalPreparedStatement statement, final String query) throws SQLException {
+        assertEquals("name=" + query, ((RelationalConnection)statement.getConnection()).getPath().getQuery());
+    }
+
+    private static void assertConnection(final MultiServerConnectionFactory.MultiServerRelationalConnection connection, final String query) {
+        assertEquals("name=" + query, connection.getPath().getQuery());
+    }
+
+    YamlRunner.YamlConnectionFactory dummyConnectionFactory(@Nonnull String name) {
         return new YamlRunner.YamlConnectionFactory() {
             @Override
             public RelationalConnection getNewConnection(@Nonnull URI connectPath) throws SQLException {
-                final RelationalConnection connection = Mockito.mock(RelationalConnection.class);
-                Mockito.when(connection.unwrap(RelationalConnection.class)).thenReturn(connection);
-                return connection;
+                // Add query string to connection so we can tell where it came from
+                URI newPath = URI.create(connectPath + "?name=" + name);
+                return dummyConnection(newPath);
             }
 
             @Override
@@ -105,4 +155,20 @@ public Set<String> getVersionsUnderTest() {
             }
         };
     }
+
+    @Nonnull
+    private static RelationalConnection dummyConnection(@Nonnull URI connectPath) throws SQLException {
+        final RelationalConnection connection = Mockito.mock(RelationalConnection.class);
+        Mockito.when(connection.unwrap(RelationalConnection.class)).thenReturn(connection);
+        Mockito.when(connection.getPath()).thenReturn(connectPath);
+        final RelationalPreparedStatement statement = dummyPreparedStatement(connection);
+        Mockito.when(connection.prepareStatement(Mockito.anyString())).thenReturn(statement);
+        return connection;
+    }
+
+    private static RelationalPreparedStatement dummyPreparedStatement(final RelationalConnection connection) throws SQLException {
+        final RelationalPreparedStatement statement = Mockito.mock(RelationalPreparedStatement.class);
+        Mockito.when(statement.getConnection()).thenReturn(connection);
+        return statement;
+    }
 }
