Minor refactor for BatchDMLE2EIT (#30364)

apache · Mar 1, 2024 · e1c8943 · e1c8943
1 parent 51e3f0a
commit e1c8943
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/...2e/sql/src/test/java/org/apache/shardingsphere/test/e2e/engine/type/dml/BaseDMLE2EIT.java b/...2e/sql/src/test/java/org/apache/shardingsphere/test/e2e/engine/type/dml/BaseDMLE2EIT.java
@@ -34,6 +34,7 @@
 import org.apache.shardingsphere.test.e2e.framework.database.DatabaseAssertionMetaData;
 import org.apache.shardingsphere.test.e2e.framework.database.DatabaseAssertionMetaDataFactory;
 import org.apache.shardingsphere.test.e2e.framework.param.model.AssertionTestParameter;
+import org.apache.shardingsphere.test.e2e.framework.param.model.CaseTestParameter;
 import org.junit.jupiter.api.AfterEach;
 
 import javax.sql.DataSource;
@@ -112,16 +113,18 @@ private void assertDataSet(final PreparedStatement actualPreparedStatement, fina
         }
     }
 
-    protected final void assertDataSet(final BatchE2EContainerComposer containerComposer, final int[] actualUpdateCounts, final DatabaseType databaseType) throws SQLException {
+    protected final void assertDataSet(final BatchE2EContainerComposer containerComposer, final int[] actualUpdateCounts, final CaseTestParameter testParam) throws SQLException {
         for (DataSetMetaData each : containerComposer.getDataSet(actualUpdateCounts).getMetaDataList()) {
-            assertDataSet(containerComposer, actualUpdateCounts, each, databaseType);
+            assertDataSet(containerComposer, actualUpdateCounts, each, testParam);
         }
     }
 
     private void assertDataSet(final BatchE2EContainerComposer containerComposer, final int[] actualUpdateCounts, final DataSetMetaData expectedDataSetMetaData,
-                               final DatabaseType databaseType) throws SQLException {
+                               final CaseTestParameter testParam) throws SQLException {
+        Map<String, DatabaseType> databaseTypes = DatabaseEnvironmentManager.getDatabaseTypes(testParam.getScenario(), testParam.getDatabaseType());
         for (String each : InlineExpressionParserFactory.newInstance(expectedDataSetMetaData.getDataNodes()).splitAndEvaluate()) {
             DataNode dataNode = new DataNode(each);
+            DatabaseType databaseType = databaseTypes.get(dataNode.getDataSourceName());
             DataSource dataSource = containerComposer.getActualDataSourceMap().get(dataNode.getDataSourceName());
             try (
                     Connection connection = dataSource.getConnection();

diff --git a/...e/sql/src/test/java/org/apache/shardingsphere/test/e2e/engine/type/dml/BatchDMLE2EIT.java b/...e/sql/src/test/java/org/apache/shardingsphere/test/e2e/engine/type/dml/BatchDMLE2EIT.java
@@ -54,7 +54,7 @@ void assertExecuteBatch(final CaseTestParameter testParam) throws SQLException,
             try (Connection connection = containerComposer.getTargetDataSource().getConnection()) {
                 actualUpdateCounts = executeBatchForPreparedStatement(testParam, connection);
             }
-            assertDataSet(containerComposer, actualUpdateCounts, testParam.getDatabaseType());
+            assertDataSet(containerComposer, actualUpdateCounts, testParam);
         }
     }