1
1
from __future__ import annotations
2
2
3
+ import importlib
3
4
import tempfile
4
5
5
6
import pytest
6
7
from luigi import LocalTarget
7
8
9
+ from gokart import file_processor
8
10
from gokart .file_processor import PolarsCsvFileProcessor , PolarsFeatherFileProcessor , PolarsJsonFileProcessor
9
11
10
12
pl = pytest .importorskip ('polars' , reason = 'polars required' )
11
13
pl_testing = pytest .importorskip ('polars.testing' , reason = 'polars required' )
12
14
13
15
14
- def test_dump_csv_with ():
16
+ @pytest .fixture
17
+ def reload_processor (monkeypatch ):
18
+ """
19
+ A pytest fixture that reloads the `gokart.file_processor` module after modifying
20
+ the environment variable `GOKART_DATAFRAME_FRAMEWORK`. This ensures that polars
21
+ is used when reloading the module.
22
+
23
+ Returns:
24
+ Tuple[Type[PolarsCsvFileProcessor], Type[PolarsFeatherFileProcessor], Type[PolarsJsonFileProcessor]]:
25
+ The reloaded classes from the `gokart.file_processor` module.
26
+ """
27
+ monkeypatch .setenv ('GOKART_DATAFRAME_FRAMEWORK' , 'polars' )
28
+ importlib .reload (file_processor )
29
+
30
+ yield PolarsCsvFileProcessor , PolarsFeatherFileProcessor , PolarsJsonFileProcessor
31
+
32
+
33
+ def test_dump_csv (reload_processor ):
34
+ PolarsCsvFileProcessor , _ , _ = reload_processor
15
35
df = pl .DataFrame ({'あ' : [1 , 2 , 3 ], 'い' : [4 , 5 , 6 ]})
16
36
processor = PolarsCsvFileProcessor ()
17
37
@@ -27,7 +47,8 @@ def test_dump_csv_with():
27
47
pl_testing .assert_frame_equal (df , loaded_df )
28
48
29
49
30
- def test_load_csv ():
50
+ def test_load_csv (reload_processor ):
51
+ PolarsCsvFileProcessor , _ , _ = reload_processor
31
52
df = pl .DataFrame ({'あ' : [1 , 2 , 3 ], 'い' : [4 , 5 , 6 ]})
32
53
processor = PolarsCsvFileProcessor ()
33
54
@@ -63,7 +84,8 @@ def test_load_csv():
63
84
pytest .param ('records' , {}, '' , id = 'With Records Orient for Empty Dict' ),
64
85
],
65
86
)
66
- def test_dump_and_load_json (orient , input_data , expected_json ):
87
+ def test_dump_and_load_json (reload_processor , orient , input_data , expected_json ):
88
+ _ , _ , PolarsJsonFileProcessor = reload_processor
67
89
processor = PolarsJsonFileProcessor (orient = orient )
68
90
69
91
with tempfile .TemporaryDirectory () as temp_dir :
@@ -82,7 +104,8 @@ def test_dump_and_load_json(orient, input_data, expected_json):
82
104
pl_testing .assert_frame_equal (df_input , loaded_df )
83
105
84
106
85
- def test_feather_should_return_same_dataframe ():
107
+ def test_feather_should_return_same_dataframe (reload_processor ):
108
+ _ , PolarsFeatherFileProcessor , _ = reload_processor
86
109
df = pl .DataFrame ({'a' : [1 ]})
87
110
# TODO: currently we set store_index_in_feather True but it is ignored
88
111
processor = PolarsFeatherFileProcessor (store_index_in_feather = True )
0 commit comments