From 5195474800bbdeb865808d431befbf916cc3645d Mon Sep 17 00:00:00 2001
From: Tom White <tom.e.white@gmail.com>
Date: Tue, 29 Oct 2024 16:02:25 +0000
Subject: [PATCH] Changes due to deprecations in Cubed

---
 tests/benchmarks/test_array.py | 6 ++----
 tests/conftest.py              | 6 +++---
 tests/utils.py                 | 4 ++--
 3 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/tests/benchmarks/test_array.py b/tests/benchmarks/test_array.py
index 4d0e62b..a58679e 100644
--- a/tests/benchmarks/test_array.py
+++ b/tests/benchmarks/test_array.py
@@ -8,7 +8,7 @@
 import cubed
 import cubed.random
 from cubed.core.optimization import multiple_inputs_optimize_dag, simple_optimize_dag
-from cubed.extensions.rich import RichProgressBar
+from cubed.diagnostics.rich import RichProgressBar
 
 from ..utils import run
 
@@ -47,9 +47,7 @@ def test_quadratic_means_xarray(tmp_path, runtime, benchmark_all, optimizer, t_l
     quad = ds**2
     quad["uv"] = ds.anom_u * ds.anom_v
     print(quad)
-    result = quad.mean(
-        "time", skipna=False, use_new_impl=True, split_every=10
-    )
+    result = quad.mean("time", skipna=False, split_every=10)
 
     if optimizer == "new-optimizer":
         opt_fn = partial(multiple_inputs_optimize_dag, max_total_num_input_blocks=20)
diff --git a/tests/conftest.py b/tests/conftest.py
index 6001cea..b528832 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -212,7 +212,7 @@ def benchmark_memory(test_run_benchmark):
     .. code-block:: python
 
         def test_something(benchmark_memory):
-            history = cubed.extensions.history.HistoryCallback()
+            history = cubed.diagnostics.history.HistoryCallback()
             with benchmark_memory(history):
                 cubed.compute(*arrs, callbacks=[history])
     """
@@ -259,7 +259,7 @@ def benchmark_tasks(test_run_benchmark):
     .. code-block:: python
 
         def test_something(benchmark_tasks):
-            history = cubed.extensions.history.HistoryCallback()
+            history = cubed.diagnostics.history.HistoryCallback()
             with benchmark_tasks(history):
                 cubed.compute(*arrs, callbacks=[history])
     """
@@ -305,7 +305,7 @@ def benchmark_all(
     .. code-block:: python
 
         def test_something(benchmark_all):
-            history = cubed.extensions.history.HistoryCallback()
+            history = cubed.diagnostics.history.HistoryCallback()
             with benchmark_all(history):
                 cubed.compute(*arrs, callbacks=[history])
     
diff --git a/tests/utils.py b/tests/utils.py
index 9257421..d365b41 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -2,8 +2,8 @@
 import yaml
 
 import cubed
-from cubed.extensions.history import HistoryCallback
-from cubed.extensions.timeline import TimelineVisualizationCallback
+from cubed.diagnostics.history import HistoryCallback
+from cubed.diagnostics.timeline import TimelineVisualizationCallback
 from cubed.spec import spec_from_config
 from cubed import config