Skip to content

Commit b48f45f

Browse files
test: Fix engine switching reset (#2729)
1 parent 0a1cc25 commit b48f45f

File tree

1 file changed

+13
-14
lines changed

1 file changed

+13
-14
lines changed

tests/unit/test_distributed.py

Lines changed: 13 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,8 @@ def wr() -> Iterator[ModuleType]:
2323
yield reload(awswrangler)
2424

2525
# Reset for future tests
26-
awswrangler.engine.register()
26+
awswrangler.engine.set(awswrangler.engine.get_installed().value)
27+
awswrangler.memory_format.set(awswrangler.memory_format.get_installed().value)
2728

2829

2930
@pytest.mark.skipif(condition=not is_ray_modin, reason="ray not available")
@@ -71,22 +72,20 @@ def test_engine_python_without_ray_installed(wr: ModuleType) -> None:
7172

7273

7374
@pytest.mark.skipif(condition=not is_ray_modin, reason="ray not available")
74-
def test_engine_switch() -> None:
75+
def test_engine_switch(wr: ModuleType) -> None:
7576
from modin.pandas import DataFrame as ModinDataFrame
7677
from pandas import DataFrame as PandasDataFrame
7778

78-
import awswrangler as wr2
79+
assert wr.engine.get_installed() == wr.EngineEnum.RAY
80+
assert wr.memory_format.get_installed() == wr.MemoryFormatEnum.MODIN
7981

80-
assert wr2.engine.get_installed() == wr2.EngineEnum.RAY
81-
assert wr2.memory_format.get_installed() == wr2.MemoryFormatEnum.MODIN
82+
assert wr.engine.get() == wr.EngineEnum.RAY
83+
assert wr.memory_format.get() == wr.MemoryFormatEnum.MODIN
84+
assert wr.pandas.DataFrame == ModinDataFrame
8285

83-
assert wr2.engine.get() == wr2.EngineEnum.RAY
84-
assert wr2.memory_format.get() == wr2.MemoryFormatEnum.MODIN
85-
assert wr2.pandas.DataFrame == ModinDataFrame
86+
wr.engine.set("python")
87+
wr.memory_format.set("pandas")
8688

87-
wr2.engine.set("python")
88-
wr2.memory_format.set("pandas")
89-
90-
assert wr2.engine.get() == wr2.EngineEnum.PYTHON
91-
assert wr2.memory_format.get() == wr2.MemoryFormatEnum.PANDAS
92-
assert wr2.pandas.DataFrame == PandasDataFrame
89+
assert wr.engine.get() == wr.EngineEnum.PYTHON
90+
assert wr.memory_format.get() == wr.MemoryFormatEnum.PANDAS
91+
assert wr.pandas.DataFrame == PandasDataFrame

0 commit comments

Comments
 (0)