Use @pytest_asyncio.fixture instead of @pytest.fixture for async fixtures

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 (

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

@@ -16,6 +16,7 @@
 import os
 
 import pytest
+import pytest_asyncio
 
 from ..... import oscar as mo
 from .....tests.core import require_ray
@@ -52,7 +53,7 @@ def index(self):
         return self._index
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def mars_cluster():
     mo.setup_cluster(address_to_resources=TEST_ADDRESS_TO_RESOURCES)
     main_pool_handles = []  # Hold actor_handle to avoid actor being freed.

mars/services/cluster/tests/test_api.py

@@ -15,6 +15,7 @@
 import asyncio
 
 import pytest
+import pytest_asyncio
 
 from .... import oscar as mo
 from ....utils import get_next_port
@@ -25,7 +26,7 @@
 from ..core import NodeStatus
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def actor_pool():
     pool = await mo.create_actor_pool("127.0.0.1", n_process=0)
     await pool.start()

mars/services/cluster/tests/test_locator.py

@@ -18,6 +18,7 @@
 from typing import List
 
 import pytest
+import pytest_asyncio
 
 from .... import oscar as mo
 from ....utils import Timer
@@ -52,7 +53,7 @@ def put_starting_nodes(self, nodes: List[str], role: NodeRole):
             self._node_infos[node] = NodeStatus.STARTING
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def actor_pool():
     pool = await mo.create_actor_pool("127.0.0.1", n_process=0)
     await pool.start()

mars/services/cluster/tests/test_procinfo.py

@@ -13,12 +13,13 @@
 # limitations under the License.
 
 import pytest
+import pytest_asyncio
 
 from .... import oscar as mo
 from ..procinfo import ProcessInfoManagerActor
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def actor_pool():
     pool = await mo.create_actor_pool(
         "127.0.0.1", n_process=2, labels=["main", "numa-0", "gpu-0"]

mars/services/cluster/tests/test_service.py

@@ -16,14 +16,15 @@
 import os
 
 import pytest
+import pytest_asyncio
 
 from .... import oscar as mo
 from ....storage import StorageLevel
 from ... import start_services, stop_services, NodeRole
 from .. import ClusterAPI, WorkerSlotInfo, QuotaInfo, StorageInfo, DiskInfo
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def actor_pools():
     async def start_pool():
         pool = await mo.create_actor_pool("127.0.0.1", n_process=0)

mars/services/cluster/tests/test_uploader.py

@@ -15,6 +15,7 @@
 import asyncio
 
 import pytest
+import pytest_asyncio
 
 from .... import oscar as mo
 from ... import NodeRole
@@ -23,7 +24,7 @@
 from ..uploader import NodeInfoUploaderActor
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def actor_pool():
     pool = await mo.create_actor_pool("127.0.0.1", n_process=0)
     await pool.start()

mars/services/mutable/tests/test_mutable.py

@@ -17,6 +17,7 @@
 import sys
 
 import pytest
+import pytest_asyncio
 import numpy as np
 
 from ....deploy.oscar.local import new_cluster
@@ -28,7 +29,7 @@
 _is_windows = sys.platform.lower().startswith("win")
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def create_cluster():
     client = await new_cluster(n_worker=2, n_cpu=2, web=True)
     async with client:

mars/services/scheduling/supervisor/tests/test_assigner.py

@@ -15,6 +15,7 @@
 import numpy as np
 import asyncio
 import pytest
+import pytest_asyncio
 
 from ..... import oscar as mo
 from .....core import ChunkGraph
@@ -100,7 +101,7 @@ async def create(cls, address: str, **kw):
         return api
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def actor_pool(request):
     pool = await mo.create_actor_pool("127.0.0.1", n_process=0)
     with_gpu = request.param

mars/services/scheduling/supervisor/tests/test_globalresource.py

@@ -15,6 +15,7 @@
 import asyncio
 
 import pytest
+import pytest_asyncio
 
 from ..... import oscar as mo
 from .....resource import Resource
@@ -23,7 +24,7 @@
 from ...supervisor import GlobalResourceManagerActor
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def actor_pool():
     pool = await mo.create_actor_pool("127.0.0.1", n_process=0)
 

mars/services/scheduling/supervisor/tests/test_manager.py

@@ -17,6 +17,7 @@
 from typing import List, Tuple, Set
 
 import pytest
+import pytest_asyncio
 
 from ..... import oscar as mo
 from .....typing import BandType
@@ -98,7 +99,7 @@ async def wait_subtask(self, subtask_id: str, band_name: str):
             pass
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def actor_pool():
     pool = await mo.create_actor_pool("127.0.0.1", n_process=0)
 

mars/services/scheduling/supervisor/tests/test_queue_balance.py

@@ -14,6 +14,7 @@
 
 import asyncio
 import pytest
+import pytest_asyncio
 from collections import defaultdict
 from typing import Tuple, List
 from ..... import oscar as mo
@@ -147,7 +148,7 @@ def dump_data(self):
         return self._submitted_subtask_ids
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def actor_pool():
     pool = await mo.create_actor_pool("127.0.0.1", n_process=0)
 

mars/services/scheduling/supervisor/tests/test_queueing.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import pytest
+import pytest_asyncio
 from typing import Tuple, List
 
 from ..... import oscar as mo
@@ -70,7 +71,7 @@ def dump_data(self):
         return self._subtask_ids, self._bands
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def actor_pool():
     pool = await mo.create_actor_pool("127.0.0.1", n_process=0)
 

mars/services/scheduling/supervisor/tests/test_speculation.py

@@ -16,6 +16,7 @@
 from typing import List, Tuple, Set
 
 import pytest
+import pytest_asyncio
 
 from ..... import oscar as mo
 from ....cluster import MockClusterAPI
@@ -54,7 +55,7 @@ async def get_exceptions(self):
         return self._exceptions
 
 
-@pytest.fixture
+@pytest_asyncio.fixture
 async def actor_pool():
     pool = await mo.create_actor_pool("127.0.0.1", n_process=0)