Fix timezone-dependent assertions in timestamp tests

Author: plusplusjiajia

paimon-python/pypaimon/tests/timestamp_test.py

@@ -15,74 +15,75 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+import calendar
 import struct
 import unittest
 from datetime import datetime
 from decimal import Decimal
 
 from pypaimon.schema.data_types import AtomicType, DataField
-from pypaimon.table.row.generic_row import GenericRow, GenericRowSerializer, GenericRowDeserializer
+from pypaimon.table.row.generic_row import (
+    GenericRow, GenericRowSerializer, GenericRowDeserializer,
+    _datetime_to_millis_and_nanos,
+)
 from pypaimon.table.row.row_kind import RowKind
 
 
+def _to_millis(dt: datetime) -> int:
+    """Helper: datetime to epoch millis using the same logic as serialization."""
+    millis, _ = _datetime_to_millis_and_nanos(dt)
+    return millis
+
+
 class TimestampTest(unittest.TestCase):
-    """Tests for timestamp serialization/deserialization in GenericRow,
-    aligned with Java BinaryRow's compact/non-compact format."""
+    """Tests for timestamp serialization/deserialization in GenericRow."""
 
     def test_timestamp_compact(self):
-        """Compact timestamp (precision <= 3): epoch millis stored directly in fixed part."""
+        """Compact timestamp (precision <= 3): round-trip preserves millis."""
         for type_str in ["TIMESTAMP(0)", "TIMESTAMP(3)"]:
             with self.subTest(type=type_str):
                 fields = [DataField(0, "ts", AtomicType(type_str))]
-                ts = datetime(2025, 4, 8, 10, 30, 0, 123000)  # .123 seconds
+                ts = datetime(2025, 4, 8, 10, 30, 0, 123000)
                 row = GenericRow([ts], fields, RowKind.INSERT)
                 serialized = GenericRowSerializer.to_bytes(row)
                 result = GenericRowDeserializer.from_bytes(serialized, fields)
-                # compact stores millis, so microsecond part is truncated to millis
-                expected_millis = int(ts.timestamp() * 1000)
-                actual_millis = int(result.values[0].timestamp() * 1000)
-                self.assertEqual(actual_millis, expected_millis)
+                self.assertEqual(_to_millis(result.values[0]), _to_millis(ts))
 
     def test_timestamp_compact_binary_format(self):
-        """Verify compact timestamp binary layout: epoch millis in fixed slot, no variable area."""
+        """Verify compact binary layout: epoch millis in fixed slot, no variable area."""
         fields = [DataField(0, "ts", AtomicType("TIMESTAMP(3)"))]
         ts = datetime(2025, 4, 8, 10, 30, 0)
         row = GenericRow([ts], fields, RowKind.INSERT)
         serialized = GenericRowSerializer.to_bytes(row)
 
-        data = serialized[4:]  # skip arity prefix
+        data = serialized[4:]
         null_bits_size = 8
         fixed_part_size = null_bits_size + 1 * 8
-        # No variable area for compact timestamp
         self.assertEqual(len(data), fixed_part_size)
-        # Fixed slot contains epoch millis
         field_offset = null_bits_size
         millis = struct.unpack('<q', data[field_offset:field_offset + 8])[0]
-        self.assertEqual(millis, int(ts.timestamp() * 1000))
+        self.assertEqual(millis, _to_millis(ts))
 
     def test_timestamp_non_compact(self):
         """Non-compact timestamp (precision 4-9): round-trip preserves microseconds."""
         for type_str in ["TIMESTAMP(6)", "TIMESTAMP(9)"]:
             with self.subTest(type=type_str):
                 fields = [DataField(0, "ts", AtomicType(type_str))]
-                ts = datetime(2025, 4, 8, 10, 30, 0, 123456)  # .123456 seconds
+                ts = datetime(2025, 4, 8, 10, 30, 0, 123456)
                 row = GenericRow([ts], fields, RowKind.INSERT)
                 serialized = GenericRowSerializer.to_bytes(row)
                 result = GenericRowDeserializer.from_bytes(serialized, fields)
-                # Python datetime has microsecond precision
                 self.assertEqual(result.values[0], ts)
 
     def test_timestamp_non_compact_binary_format(self):
-        """Verify non-compact timestamp binary layout matches Java:
-        - fixed slot: (offset << 32) | nanoOfMillisecond
-        - variable area: epoch millis (8 bytes)
-        """
+        """Verify non-compact binary layout: (offset << 32 | nanoOfMilli) in fixed slot,
+        epoch millis in variable area."""
         fields = [DataField(0, "ts", AtomicType("TIMESTAMP(6)"))]
-        ts = datetime(2025, 4, 8, 10, 30, 0, 123456)  # .123456 seconds
+        ts = datetime(2025, 4, 8, 10, 30, 0, 123456)
         row = GenericRow([ts], fields, RowKind.INSERT)
         serialized = GenericRowSerializer.to_bytes(row)
 
-        data = serialized[4:]  # skip arity prefix
+        data = serialized[4:]
         null_bits_size = 8
         fixed_part_size = null_bits_size + 1 * 8
 
@@ -95,15 +96,13 @@ def test_timestamp_non_compact_binary_format(self):
         cursor = (offset_and_nano >> 32) & 0xFFFFFFFF
         nano_of_millisecond = offset_and_nano & 0xFFFFFFFF
 
-        # cursor should point to variable area
         self.assertEqual(cursor, fixed_part_size)
         # 123456 us = 123 ms + 456 us = 123 ms + 456000 ns
         self.assertEqual(nano_of_millisecond, 456000)
 
         # Variable area contains epoch millis
         var_millis = struct.unpack('<q', data[cursor:cursor + 8])[0]
-        expected_millis = int(ts.timestamp() * 1000)
-        self.assertEqual(var_millis, expected_millis)
+        self.assertEqual(var_millis, _to_millis(ts))
 
     def test_timestamp_non_compact_null(self):
         """Non-compact timestamp null value."""
@@ -116,19 +115,17 @@ def test_timestamp_non_compact_null(self):
     def test_timestamp_boundary_precision(self):
         """Precision 3 is last compact, precision 4 is first non-compact."""
         ts = datetime(2025, 4, 8, 10, 30, 0, 123456)
+        fixed_part_size = 8 + 1 * 8
 
         # precision=3: compact, no variable area
         fields_3 = [DataField(0, "ts", AtomicType("TIMESTAMP(3)"))]
         s_3 = GenericRowSerializer.to_bytes(GenericRow([ts], fields_3, RowKind.INSERT))
-        data_3 = s_3[4:]
-        fixed_part_size = 8 + 1 * 8
-        self.assertEqual(len(data_3), fixed_part_size)
+        self.assertEqual(len(s_3[4:]), fixed_part_size)
 
         # precision=4: non-compact, has 8-byte variable area
         fields_4 = [DataField(0, "ts", AtomicType("TIMESTAMP(4)"))]
         s_4 = GenericRowSerializer.to_bytes(GenericRow([ts], fields_4, RowKind.INSERT))
-        data_4 = s_4[4:]
-        self.assertEqual(len(data_4), fixed_part_size + 8)
+        self.assertEqual(len(s_4[4:]), fixed_part_size + 8)
 
     def test_timestamp_mixed_with_other_types(self):
         """Non-compact timestamp mixed with other types in a single row."""
@@ -153,7 +150,7 @@ def test_timestamp_mixed_with_other_types(self):
 
         self.assertEqual(result.values[0], 42)
         self.assertEqual(result.values[1], "hello")
-        self.assertEqual(int(result.values[2].timestamp() * 1000), int(ts_compact.timestamp() * 1000))
+        self.assertEqual(_to_millis(result.values[2]), _to_millis(ts_compact))
         self.assertEqual(result.values[3], ts_non_compact)
         self.assertEqual(result.values[4], dec_val)
 
@@ -164,15 +161,13 @@ def test_timestamp_default_precision(self):
         row = GenericRow([ts], fields, RowKind.INSERT)
         serialized = GenericRowSerializer.to_bytes(row)
         result = GenericRowDeserializer.from_bytes(serialized, fields)
-        expected_millis = int(ts.timestamp() * 1000)
-        actual_millis = int(result.values[0].timestamp() * 1000)
-        self.assertEqual(actual_millis, expected_millis)
+        self.assertEqual(_to_millis(result.values[0]), _to_millis(ts))
 
     def test_timestamp_pre_epoch(self):
         """Dates before 1970-01-01 (negative epoch millis) should round-trip correctly."""
         # Compact
         fields3 = [DataField(0, "ts", AtomicType("TIMESTAMP(3)"))]
-        ts_pre = datetime(1969, 7, 20, 20, 17, 0)  # Apollo 11 moon landing
+        ts_pre = datetime(1969, 7, 20, 20, 17, 0)
         row = GenericRow([ts_pre], fields3, RowKind.INSERT)
         serialized = GenericRowSerializer.to_bytes(row)
         result = GenericRowDeserializer.from_bytes(serialized, fields3)