Use @pytest_asyncio.fixture instead of @pytest.fixture for async fixtures (#3025)

Author: fyrestone

mars/conftest.py

@@ -18,6 +18,7 @@
 
 import psutil
 import pytest
+import pytest_asyncio
 
 from mars.config import option_context
 from mars.core.mode import is_kernel_mode, is_build_mode
@@ -132,7 +133,7 @@ def stop_ray(request):  # pragma: no cover
         ray.shutdown()
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def ray_create_mars_cluster(request):
     from mars.deploy.oscar.ray import new_cluster, _load_config
 

mars/dataframe/contrib/raydataset/tests/test_raydataset.py

@@ -16,6 +16,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+import pytest_asyncio
 
 from ..... import dataframe as md
 from .....deploy.oscar.ray import new_cluster
@@ -39,7 +40,7 @@
     sklearn = None
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def create_cluster(request):
     client = await new_cluster(
         supervisor_mem=256 * 1024**2,

mars/deploy/oscar/tests/test_fault_injection.py

@@ -14,6 +14,7 @@
 
 import os
 import pytest
+import pytest_asyncio
 import traceback
 import numpy as np
 import pandas as pd
@@ -40,7 +41,7 @@
 )
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def fault_cluster(request):
     param = getattr(request, "param", {})
     start_method = os.environ.get("POOL_START_METHOD", None)

mars/deploy/oscar/tests/test_local.py

@@ -23,6 +23,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+import pytest_asyncio
 
 try:
     import vineyard
@@ -103,7 +104,7 @@
 
 
 @pytest.mark.parametrize(indirect=True)
-@pytest.fixture(params=params)
+@pytest_asyncio.fixture(params=params)
 async def create_cluster(request):
     if request.param == "default":
         config = CONFIG_TEST_FILE
@@ -822,7 +823,7 @@ async def _exec():
 min_task_runtime = 2
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def speculative_cluster():
     config = _load_config()
     config["scheduling"]["speculation"]["enabled"] = True

mars/deploy/oscar/tests/test_ray.py

@@ -24,6 +24,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+import pytest_asyncio
 
 import mars
 from .... import oscar as mo
@@ -90,7 +91,7 @@
 EXPECT_PROFILING_STRUCTURE_NO_SLOW["supervisor"]["slow_subtasks"] = {}
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def create_cluster(request):
     param = getattr(request, "param", {})
     ray_config = _load_config(CONFIG_FILE)

mars/deploy/oscar/tests/test_ray_fault_injection.py

@@ -15,6 +15,7 @@
 import os
 
 import pytest
+import pytest_asyncio
 
 from ....oscar.errors import ServerClosed
 from ....services.tests.fault_injection_manager import (
@@ -45,7 +46,7 @@
 }
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def fault_cluster(request):
     param = getattr(request, "param", {})
     ray_config = _load_config(RAY_CONFIG_FILE)

mars/deploy/oscar/tests/test_ray_scheduling.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import pytest
+import pytest_asyncio
 
 from ....tests.core import require_ray
 from ....utils import lazy_import
@@ -26,7 +27,7 @@
 ray = lazy_import("ray")
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def speculative_cluster():
     client = await new_cluster(
         "test_cluster",

mars/oscar/backends/mars/tests/test_debug.py

@@ -21,6 +21,7 @@
 from typing import List
 
 import pytest
+import pytest_asyncio
 
 from ..... import oscar as mo
 from ....debug import reload_debug_opts_from_env, get_debug_options
@@ -62,7 +63,7 @@ async def call_self_ref(self):
         await self.ref().wait(1)
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def actor_pool():
     start_method = (
         os.environ.get("POOL_START_METHOD", "forkserver")

mars/oscar/backends/mars/tests/test_mars_actor_context.py

@@ -22,6 +22,7 @@
 
 import pandas as pd
 import pytest
+import pytest_asyncio
 
 from ..... import oscar as mo
 from .....oscar.core import ActorRef, LocalActorRef
@@ -241,7 +242,7 @@ def get_call_log(self):
 
 
 @pytest.mark.parametrize(indirect=True)
-@pytest.fixture(params=[False, True])
+@pytest_asyncio.fixture(params=[False, True])
 async def actor_pool(request):
     start_method = (
         os.environ.get("POOL_START_METHOD", "forkserver")

mars/oscar/backends/ray/tests/test_ray_actor_context.py

@@ -16,6 +16,7 @@
 import time
 
 import pytest
+import pytest_asyncio
 
 from .....utils import lazy_import
 from .....tests.core import require_ray
@@ -29,7 +30,7 @@
 ray = lazy_import("ray")
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def actor_pool_context():
     pg_name, n_process = f"ray_cluster_{time.time_ns()}", 2
     from .....serialization.ray import (