From da8657e994bfad2214eeeebb3c50f4e46d86577a Mon Sep 17 00:00:00 2001
From: Tom White <tom.e.white@gmail.com>
Date: Wed, 14 Aug 2024 14:22:34 +0100
Subject: [PATCH] Add a --use-cubed option to pytest

---
 conftest.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/conftest.py b/conftest.py
index bcccb6e8a..9343393f3 100644
--- a/conftest.py
+++ b/conftest.py
@@ -2,6 +2,30 @@
 collect_ignore_glob = ["benchmarks/**", "sgkit/io/vcf/*.py", ".github/scripts/*.py"]
 
 
+def pytest_addoption(parser):
+    parser.addoption(
+        "--use-cubed", action="store_true", default=False, help="run with cubed"
+    )
+
+
+def use_cubed():
+    import dask
+    import xarray as xr
+
+    # set xarray to use cubed by default
+    xr.set_options(chunk_manager="cubed")
+
+    # ensure that dask compute raises if it is ever called
+    class AlwaysRaiseScheduler:
+        def __call__(self, dsk, keys, **kwargs):
+            raise RuntimeError("Dask 'compute' was called")
+
+    dask.config.set(scheduler=AlwaysRaiseScheduler())
+
+
 def pytest_configure(config) -> None:  # type: ignore
     # Add "gpu" marker
     config.addinivalue_line("markers", "gpu:Run tests that run on GPU")
+
+    if config.getoption("--use-cubed"):
+        use_cubed()