diff --git a/python/cudf/cudf/tests/test_pack.py b/python/cudf/cudf/tests/test_pack.py index 9972071122e..9b5a8c19cf5 100644 --- a/python/cudf/cudf/tests/test_pack.py +++ b/python/cudf/cudf/tests/test_pack.py @@ -191,15 +191,14 @@ def check_packed_pickled_equality(df): assert isinstance(sortvaldf.index, GenericIndex) assert_packed_frame_picklable(sortvaldf) # out-of-band - if pickle.HIGHEST_PROTOCOL >= 5: - buffers = [] - serialbytes = pickle.dumps( - pack(df), protocol=5, buffer_callback=buffers.append - ) - for b in buffers: - assert isinstance(b, pickle.PickleBuffer) - loaded = unpack(pickle.loads(serialbytes, buffers=buffers)) - assert_eq(loaded, df) + buffers = [] + serialbytes = pickle.dumps( + pack(df), protocol=5, buffer_callback=buffers.append + ) + for b in buffers: + assert isinstance(b, pickle.PickleBuffer) + loaded = unpack(pickle.loads(serialbytes, buffers=buffers)) + assert_eq(loaded, df) def assert_packed_frame_picklable(df): diff --git a/python/cudf/cudf/tests/test_pickling.py b/python/cudf/cudf/tests/test_pickling.py index 8ce818e7a3d..71c1f206a64 100644 --- a/python/cudf/cudf/tests/test_pickling.py +++ b/python/cudf/cudf/tests/test_pickling.py @@ -1,4 +1,4 @@ -# Copyright (c) 2018-2022, NVIDIA CORPORATION. +# Copyright (c) 2018-2023, NVIDIA CORPORATION. import pickle @@ -25,15 +25,12 @@ def check_serialization(df): assert isinstance(sortvaldf.index, (GenericIndex, RangeIndex)) assert_frame_picklable(sortvaldf) # out-of-band - if pickle.HIGHEST_PROTOCOL >= 5: - buffers = [] - serialbytes = pickle.dumps( - df, protocol=5, buffer_callback=buffers.append - ) - for b in buffers: - assert isinstance(b, pickle.PickleBuffer) - loaded = pickle.loads(serialbytes, buffers=buffers) - assert_eq(loaded, df) + buffers = [] + serialbytes = pickle.dumps(df, protocol=5, buffer_callback=buffers.append) + for b in buffers: + assert isinstance(b, pickle.PickleBuffer) + loaded = pickle.loads(serialbytes, buffers=buffers) + assert_eq(loaded, df) def assert_frame_picklable(df):