split methods

Author: Ilyas Gasanov

tests/test_integration/test_run_transfer/test_clickhouse.py

@@ -13,7 +13,7 @@
 from tests.mocks import MockUser
 from tests.test_unit.utils import create_transfer
 from tests.utils import (
-    prepare_dataframes_for_comparison,
+    cast_dataframe_types,
     run_transfer_and_verify,
     split_df,
 )
@@ -131,7 +131,7 @@ async def test_run_transfer_postgres_to_clickhouse_with_full_strategy(
     )
     df = reader.run()
 
-    df, init_df = prepare_dataframes_for_comparison(df, init_df)
+    df, init_df = cast_dataframe_types(df, init_df)
     assert df.sort("ID").collect() == init_df.sort("ID").collect()
 
 
@@ -169,7 +169,7 @@ async def test_run_transfer_postgres_to_clickhouse_mixed_naming_with_full_strate
     assert df.columns != init_df_with_mixed_column_naming.columns
     assert df.columns == [column.lower() for column in init_df_with_mixed_column_naming.columns]
 
-    df, init_df_with_mixed_column_naming = prepare_dataframes_for_comparison(df, init_df_with_mixed_column_naming)
+    df, init_df_with_mixed_column_naming = cast_dataframe_types(df, init_df_with_mixed_column_naming)
     assert df.sort("ID").collect() == init_df_with_mixed_column_naming.sort("ID").collect()
 
 
@@ -205,14 +205,14 @@ async def test_run_transfer_postgres_to_clickhouse_with_incremental_strategy(
     )
     df = reader.run()
 
-    df, first_transfer_df = prepare_dataframes_for_comparison(df, first_transfer_df)
+    df, first_transfer_df = cast_dataframe_types(df, first_transfer_df)
     assert df.sort("ID").collect() == first_transfer_df.sort("ID").collect()
 
     fill_with_data(second_transfer_df)
     await run_transfer_and_verify(client, group_owner, postgres_to_clickhouse.id)
 
     df_with_increment = reader.run()
-    df_with_increment, init_df = prepare_dataframes_for_comparison(df_with_increment, init_df)
+    df_with_increment, init_df = cast_dataframe_types(df_with_increment, init_df)
     assert df_with_increment.sort("ID").collect() == init_df.sort("ID").collect()
 
 
@@ -258,7 +258,7 @@ async def test_run_transfer_clickhouse_to_postgres_with_full_strategy(
     )
     df = reader.run()
 
-    df, init_df = prepare_dataframes_for_comparison(df, init_df)
+    df, init_df = cast_dataframe_types(df, init_df)
     assert df.sort("ID").collect() == init_df.sort("ID").collect()
 
 
@@ -296,7 +296,7 @@ async def test_run_transfer_clickhouse_to_postgres_mixed_naming_with_full_strate
     assert df.columns != init_df_with_mixed_column_naming.columns
     assert df.columns == [column.lower() for column in init_df_with_mixed_column_naming.columns]
 
-    df, init_df_with_mixed_column_naming = prepare_dataframes_for_comparison(df, init_df_with_mixed_column_naming)
+    df, init_df_with_mixed_column_naming = cast_dataframe_types(df, init_df_with_mixed_column_naming)
     assert df.sort("ID").collect() == init_df_with_mixed_column_naming.sort("ID").collect()
 
 
@@ -332,12 +332,12 @@ async def test_run_transfer_clickhouse_to_postgres_with_incremental_strategy(
     )
     df = reader.run()
 
-    df, first_transfer_df = prepare_dataframes_for_comparison(df, first_transfer_df)
+    df, first_transfer_df = cast_dataframe_types(df, first_transfer_df)
     assert df.sort("ID").collect() == first_transfer_df.sort("ID").collect()
 
     fill_with_data(second_transfer_df)
     await run_transfer_and_verify(client, group_owner, clickhouse_to_postgres.id)
 
     df_with_increment = reader.run()
-    df_with_increment, init_df = prepare_dataframes_for_comparison(df_with_increment, init_df)
+    df_with_increment, init_df = cast_dataframe_types(df_with_increment, init_df)
     assert df_with_increment.sort("ID").collect() == init_df.sort("ID").collect()

tests/test_integration/test_run_transfer/test_ftp.py

@@ -18,9 +18,10 @@
 from tests.test_unit.utils import create_transfer
 from tests.utils import (
     add_increment_to_files_and_upload,
-    prepare_dataframes_for_comparison,
+    cast_dataframe_types,
     run_transfer_and_verify,
     split_df,
+    truncate_datetime_to_seconds,
     verify_file_name_template,
 )
 
@@ -181,12 +182,10 @@ async def test_run_transfer_ftp_to_postgres_with_full_strategy(
     )
     df = reader.run()
 
-    df, init_df = prepare_dataframes_for_comparison(
-        df,
-        init_df,
-        file_format=file_format,
-        transfer_direction="file_to_db",
-    )
+    if file_format == "excel":
+        df, init_df = truncate_datetime_to_seconds(df, init_df, transfer_direction="file_to_db")
+
+    df, init_df = cast_dataframe_types(df, init_df)
     assert df.sort("id").collect() == init_df.sort("id").collect()
 
 
@@ -225,12 +224,7 @@ async def test_run_transfer_ftp_to_postgres_with_incremental_strategy(
     )
     df = reader.run()
 
-    df, init_df = prepare_dataframes_for_comparison(
-        df,
-        init_df,
-        file_format=file_format,
-        transfer_direction="file_to_db",
-    )
+    df, init_df = cast_dataframe_types(df, init_df)
     assert df.sort("id").collect() == init_df.sort("id").collect()
 
     add_increment_to_files_and_upload(
@@ -242,12 +236,7 @@ async def test_run_transfer_ftp_to_postgres_with_incremental_strategy(
     await run_transfer_and_verify(client, group_owner, ftp_to_postgres.id)
 
     df_with_increment = reader.run()
-    df_with_increment, init_df = prepare_dataframes_for_comparison(
-        df_with_increment,
-        init_df,
-        file_format=file_format,
-        transfer_direction="file_to_db",
-    )
+    df_with_increment, init_df = cast_dataframe_types(df_with_increment, init_df)
     assert df_with_increment.sort("id").collect() == init_df.union(init_df).sort("id").collect()
 
 
@@ -337,12 +326,10 @@ async def test_run_transfer_postgres_to_ftp_with_full_strategy(
     )
     df = reader.run()
 
-    df, init_df = prepare_dataframes_for_comparison(
-        df,
-        init_df,
-        file_format=format_name,
-        transfer_direction="db_to_file",
-    )
+    if format_name == "excel":
+        df, init_df = truncate_datetime_to_seconds(df, init_df, transfer_direction="db_to_file")
+
+    df, init_df = cast_dataframe_types(df, init_df)
     assert df.sort("id").collect() == init_df.sort("id").collect()
 
 
@@ -399,12 +386,7 @@ async def test_run_transfer_postgres_to_ftp_with_incremental_strategy(
     )
     df = reader.run()
 
-    df, first_transfer_df = prepare_dataframes_for_comparison(
-        df,
-        first_transfer_df,
-        file_format=format_name,
-        transfer_direction="db_to_file",
-    )
+    df, first_transfer_df = cast_dataframe_types(df, first_transfer_df)
     assert df.sort("id").collect() == first_transfer_df.sort("id").collect()
 
     fill_with_data(second_transfer_df)
@@ -414,10 +396,5 @@ async def test_run_transfer_postgres_to_ftp_with_incremental_strategy(
     verify_file_name_template(os.listdir(tmp_path), expected_extension)
 
     df_with_increment = reader.run()
-    df_with_increment, second_transfer_df = prepare_dataframes_for_comparison(
-        df_with_increment,
-        second_transfer_df,
-        file_format=format_name,
-        transfer_direction="db_to_file",
-    )
+    df_with_increment, second_transfer_df = cast_dataframe_types(df_with_increment, second_transfer_df)
     assert df_with_increment.sort("id").collect() == init_df.sort("id").collect()

tests/test_integration/test_run_transfer/test_ftps.py

@@ -18,7 +18,7 @@
 from tests.test_unit.utils import create_transfer
 from tests.utils import (
     add_increment_to_files_and_upload,
-    prepare_dataframes_for_comparison,
+    cast_dataframe_types,
     run_transfer_and_verify,
     split_df,
     verify_file_name_template,
@@ -145,12 +145,7 @@ async def test_run_transfer_ftps_to_postgres_with_full_strategy(
     )
     df = reader.run()
 
-    df, init_df = prepare_dataframes_for_comparison(
-        df,
-        init_df,
-        file_format=file_format,
-        transfer_direction="file_to_db",
-    )
+    df, init_df = cast_dataframe_types(df, init_df)
     assert df.sort("id").collect() == init_df.sort("id").collect()
 
 
@@ -189,12 +184,7 @@ async def test_run_transfer_ftps_to_postgres_with_incremental_strategy(
     )
     df = reader.run()
 
-    df, init_df = prepare_dataframes_for_comparison(
-        df,
-        init_df,
-        file_format=file_format,
-        transfer_direction="file_to_db",
-    )
+    df, init_df = cast_dataframe_types(df, init_df)
     assert df.sort("id").collect() == init_df.sort("id").collect()
 
     add_increment_to_files_and_upload(
@@ -206,12 +196,7 @@ async def test_run_transfer_ftps_to_postgres_with_incremental_strategy(
     await run_transfer_and_verify(client, group_owner, ftps_to_postgres.id)
 
     df_with_increment = reader.run()
-    df_with_increment, init_df = prepare_dataframes_for_comparison(
-        df_with_increment,
-        init_df,
-        file_format=file_format,
-        transfer_direction="file_to_db",
-    )
+    df_with_increment, init_df = cast_dataframe_types(df_with_increment, init_df)
     assert df_with_increment.sort("id").collect() == init_df.union(init_df).sort("id").collect()
 
 
@@ -266,12 +251,7 @@ async def test_run_transfer_postgres_to_ftps_with_full_strategy(
     )
     df = reader.run()
 
-    df, init_df = prepare_dataframes_for_comparison(
-        df,
-        init_df,
-        file_format=format_name,
-        transfer_direction="db_to_file",
-    )
+    df, init_df = cast_dataframe_types(df, init_df)
     assert df.sort("id").collect() == init_df.sort("id").collect()
 
 
@@ -328,12 +308,7 @@ async def test_run_transfer_postgres_to_ftps_with_incremental_strategy(
     )
     df = reader.run()
 
-    df, first_transfer_df = prepare_dataframes_for_comparison(
-        df,
-        first_transfer_df,
-        file_format=format_name,
-        transfer_direction="db_to_file",
-    )
+    df, first_transfer_df = cast_dataframe_types(df, first_transfer_df)
     assert df.sort("id").collect() == first_transfer_df.sort("id").collect()
 
     fill_with_data(second_transfer_df)
@@ -343,10 +318,5 @@ async def test_run_transfer_postgres_to_ftps_with_incremental_strategy(
     verify_file_name_template(os.listdir(tmp_path), expected_extension)
 
     df_with_increment = reader.run()
-    df_with_increment, second_transfer_df = prepare_dataframes_for_comparison(
-        df_with_increment,
-        second_transfer_df,
-        file_format=format_name,
-        transfer_direction="db_to_file",
-    )
+    df, second_transfer_df = cast_dataframe_types(df, second_transfer_df)
     assert df_with_increment.sort("id").collect() == init_df.sort("id").collect()

tests/test_integration/test_run_transfer/test_hdfs.py

@@ -16,9 +16,10 @@
 from tests.mocks import MockUser
 from tests.test_unit.utils import create_transfer
 from tests.utils import (
-    prepare_dataframes_for_comparison,
+    cast_dataframe_types,
     run_transfer_and_verify,
     split_df,
+    truncate_datetime_to_seconds,
     verify_file_name_template,
 )
 
@@ -169,12 +170,10 @@ async def test_run_transfer_hdfs_to_postgres_with_full_strategy(
     )
     df = reader.run()
 
-    df, init_df = prepare_dataframes_for_comparison(
-        df,
-        init_df,
-        file_format=file_format,
-        transfer_direction="file_to_db",
-    )
+    if file_format == "excel":
+        df, init_df = truncate_datetime_to_seconds(df, init_df, transfer_direction="file_to_db")
+
+    df, init_df = cast_dataframe_types(df, init_df)
     assert df.sort("id").collect() == init_df.sort("id").collect()
 
 
@@ -259,12 +258,10 @@ async def test_run_transfer_postgres_to_hdfs_with_full_strategy(
     )
     df = reader.run()
 
-    df, init_df = prepare_dataframes_for_comparison(
-        df,
-        init_df,
-        file_format=format_name,
-        transfer_direction="db_to_file",
-    )
+    if format_name == "excel":
+        df, init_df = truncate_datetime_to_seconds(df, init_df, transfer_direction="db_to_file")
+
+    df, init_df = cast_dataframe_types(df, init_df)
     assert df.sort("id").collect() == init_df.sort("id").collect()
 
 
@@ -315,12 +312,7 @@ async def test_run_transfer_postgres_to_hdfs_with_incremental_strategy(
     )
     df = reader.run()
 
-    df, first_transfer_df = prepare_dataframes_for_comparison(
-        df,
-        first_transfer_df,
-        file_format=format_name,
-        transfer_direction="db_to_file",
-    )
+    df, first_transfer_df = cast_dataframe_types(df, first_transfer_df)
     assert df.sort("id").collect() == first_transfer_df.sort("id").collect()
 
     fill_with_data(second_transfer_df)
@@ -330,10 +322,5 @@ async def test_run_transfer_postgres_to_hdfs_with_incremental_strategy(
     verify_file_name_template(files, expected_extension)
 
     df_with_increment = reader.run()
-    df_with_increment, init_df = prepare_dataframes_for_comparison(
-        df_with_increment,
-        init_df,
-        file_format=format_name,
-        transfer_direction="db_to_file",
-    )
+    df_with_increment, init_df = cast_dataframe_types(df_with_increment, init_df)
     assert df_with_increment.sort("id").collect() == init_df.sort("id").collect()