Fix xlang test.

Author: claudevdm

sdks/python/apache_beam/io/external/xlang_jdbcio_it_test.py

@@ -244,7 +244,6 @@ def test_xlang_jdbc_read_with_explicit_schema(self, database):
     else:
       binary_type = ('BINARY(10)', 'VARBINARY(10)')
 
-    # Create a test table
     with self.engine.begin() as connection:
       connection.execute(
           sqlalchemy.text(
@@ -255,7 +254,6 @@ def test_xlang_jdbc_read_with_explicit_schema(self, database):
               "f_timestamp TIMESTAMP(3), " + "f_decimal DECIMAL(10, 2), " +
               "f_date DATE, " + "f_time TIME(3))"))
 
-    # Insert test data
     inserted_rows = [
         JdbcTestRow(
             i,
@@ -271,63 +269,54 @@ def test_xlang_jdbc_read_with_explicit_schema(self, database):
         for i in range(ROW_COUNT)
     ]
 
-    # Insert the data using SQLAlchemy
-    with self.engine.begin() as connection:
-      for row in inserted_rows:
-        connection.execute(
-            sqlalchemy.text(
-                f"INSERT INTO {table_name} "
-                f"VALUES (:f_id, :f_float, :f_char, :f_varchar, "
-                f":f_bytes, :f_varbytes, :f_timestamp, :f_decimal, "
-                f":f_date, :f_time)"),
-            {
-                "f_id": row.f_id,
-                "f_float": row.f_float,
-                "f_char": row.f_char,
-                "f_varchar": row.f_varchar,
-                "f_bytes": row.f_bytes,
-                "f_varbytes": row.f_varbytes,
-                "f_timestamp": row.f_timestamp.to_utc_datetime(),
-                "f_decimal": row.f_decimal,
-                "f_date": row.f_date,
-                "f_time": row.f_time,
-            })
+    with TestPipeline() as p:
+      p.not_use_test_runner_api = True
+      _ = (
+          p
+          | beam.Create(inserted_rows).with_output_types(JdbcTestRow)
+          | 'Write to jdbc' >> WriteToJdbc(
+              table_name=table_name,
+              driver_class_name=self.driver_class_name,
+              jdbc_url=self.jdbc_url,
+              username=self.username,
+              password=self.password,
+              classpath=classpath,
+          ))
 
     # Define a custom schema with different field names
     CustomSchemaRow = typing.NamedTuple(
         "CustomSchemaRow",
         [
-            ("custom_id", int),
-            ("custom_float", float),
-            ("custom_char", str),
-            ("custom_varchar", str),
-            ("custom_bytes", bytes),
-            ("custom_varbytes", bytes),
-            ("custom_timestamp", Timestamp),
-            ("custom_decimal", Decimal),
-            ("custom_date", datetime.date),
-            ("custom_time", datetime.time),
+            ("renamed_id", int),
+            ("renamed_float", float),
+            ("renamed_char", str),
+            ("renamed_varchar", str),
+            ("renamed_bytes", bytes),
+            ("renamed_varbytes", bytes),
+            ("renamed_timestamp", Timestamp),
+            ("renamed_decimal", Decimal),
+            ("renamed_date", datetime.date),
+            ("renamed_time", datetime.time),
         ],
     )
     coders.registry.register_coder(CustomSchemaRow, coders.RowCoder)
 
     # Register MillisInstant logical type to override the mapping from Timestamp
     LogicalType.register_logical_type(MillisInstant)
 
-    # Expected results with custom field names
     expected_row = []
     for row in inserted_rows:
-      f_char = row.f_char + ' ' * (10 - len(row.f_char))
       if database != 'postgres':
-        # padding expected results
+        # padding expected results for binary fields
         f_bytes = row.f_bytes + b'\0' * (10 - len(row.f_bytes))
       else:
         f_bytes = row.f_bytes
+
       expected_row.append(
           CustomSchemaRow(
               row.f_id,
               row.f_float,
-              f_char,
+              row.f_char + ' ' * (10 - len(row.f_char)),
               row.f_varchar,
               f_bytes,
               row.f_bytes,
@@ -336,6 +325,28 @@ def test_xlang_jdbc_read_with_explicit_schema(self, database):
               row.f_date,
               row.f_time))
 
+    # Custom equals function that verifies field names and handles padding
+    def custom_equals(expected, actual):
+
+      # Then, compare the field values
+      return (
+          expected.renamed_id == actual.renamed_id and
+          expected.renamed_float == actual.renamed_float and
+          expected.renamed_char == actual.renamed_char and
+          expected.renamed_varchar == actual.renamed_varchar and (
+              expected.renamed_bytes == actual.renamed_bytes or
+              # Handle potential padding differences in binary fields
+              expected.renamed_bytes.rstrip(b'\0') ==
+              actual.renamed_bytes.rstrip(b'\0')) and (
+                  expected.renamed_varbytes == actual.renamed_varbytes or
+                  # Handle potential padding differences in binary fields
+                  expected.renamed_varbytes.rstrip(b'\0') ==
+                  actual.renamed_varbytes.rstrip(b'\0')) and
+          expected.renamed_timestamp == actual.renamed_timestamp and
+          expected.renamed_decimal == actual.renamed_decimal and
+          expected.renamed_date == actual.renamed_date and
+          expected.renamed_time == actual.renamed_time)
+
     with TestPipeline() as p:
       p.not_use_test_runner_api = True
       result = (
@@ -349,7 +360,8 @@ def test_xlang_jdbc_read_with_explicit_schema(self, database):
               classpath=classpath,
               schema=CustomSchemaRow))
 
-      assert_that(result, equal_to(expected_row))
+      # Use our custom equals function that verifies field names
+      assert_that(result, equal_to(expected_row, equals_fn=custom_equals))
 
   # Creating a container with testcontainers sometimes raises ReadTimeout
   # error. In java there are 2 retries set by default.