Wip update tests

Author: jenshnielsen

src/qcodes/dataset/exporters/export_to_xarray.py

@@ -248,11 +248,10 @@ def load_to_xarray_dataset_dict(
         dataset, datadict, use_multi_index=use_multi_index
     )
 
-    for dataname, xr_dataset in xr_datasets.items():
-        _add_param_spec_to_xarray_coords(dataset, xr_dataset[dataname])
-        paramspec_dict = _paramspec_dict_with_extras(dataset, str(dataname))
-        xr_dataset[dataname].attrs.update(paramspec_dict.items())
-        _add_metadata_to_xarray(dataset, xr_dataset[dataname])
+    for xr_dataset in xr_datasets.values():
+        _add_param_spec_to_xarray_coords(dataset, xr_dataset)
+        _add_param_spec_to_xarray_data_vars(dataset, xr_dataset)
+        _add_metadata_to_xarray(dataset, xr_dataset)
 
     return xr_datasets
 

tests/dataset/test_dataset_export.py

@@ -40,6 +40,7 @@
 from qcodes.dataset.exporters.export_to_xarray import _calculate_index_shape
 from qcodes.dataset.linked_datasets.links import links_to_str
 from qcodes.parameters import ManualParameter, Parameter
+from qcodes.utils.deprecate import QCoDeSDeprecationWarning
 
 if TYPE_CHECKING:
     from collections.abc import Hashable
@@ -868,7 +869,17 @@ def test_export_to_xarray_dataset_empty_ds(mock_empty_dataset: DataSet) -> None:
 
 
 def test_export_to_xarray_dataarray_empty_ds(mock_empty_dataset: DataSet) -> None:
-    dad = mock_empty_dataset.to_xarray_dataarray_dict()
+    with pytest.warns(QCoDeSDeprecationWarning, match="to_xarray_dataarray_dict"):
+        dad = mock_empty_dataset.to_xarray_dataarray_dict()  # pyright: ignore[reportDeprecated]
+    assert len(dad) == 2
+    assert len(dad["y"].coords) == 1
+    assert "x" in dad["y"].coords
+    assert len(dad["z"].coords) == 1
+    assert "x" in dad["z"].coords
+
+
+def test_export_to_xarray_dataset_dict_empty_ds(mock_empty_dataset: DataSet) -> None:
+    dad = mock_empty_dataset.to_xarray_dataset_dict()
     assert len(dad) == 2
     assert len(dad["y"].coords) == 1
     assert "x" in dad["y"].coords
@@ -910,10 +921,20 @@ def test_export_to_xarray_extra_metadata(mock_dataset: DataSet) -> None:
         assert "snapshot" not in ds[array_name].attrs.keys()
 
 
+def test_export_to_xarray_da_dict_extra_metadata(mock_dataset: DataSet) -> None:
+    mock_dataset.add_metadata("mytag", "somestring")
+    mock_dataset.add_metadata("myothertag", 1)
+    with pytest.warns(QCoDeSDeprecationWarning, match="to_xarray_dataarray_dict"):
+        da_dict = mock_dataset.to_xarray_dataarray_dict()  # pyright: ignore[reportDeprecated]
+
+    for datarray in da_dict.values():
+        _assert_xarray_metadata_is_as_expected(datarray, mock_dataset)
+
+
 def test_export_to_xarray_ds_dict_extra_metadata(mock_dataset: DataSet) -> None:
     mock_dataset.add_metadata("mytag", "somestring")
     mock_dataset.add_metadata("myothertag", 1)
-    da_dict = mock_dataset.to_xarray_dataarray_dict()
+    da_dict = mock_dataset.to_xarray_dataset_dict()
 
     for datarray in da_dict.values():
         _assert_xarray_metadata_is_as_expected(datarray, mock_dataset)
@@ -973,7 +994,8 @@ def test_to_xarray_ds_paramspec_metadata_is_preserved(
 def test_to_xarray_da_dict_paramspec_metadata_is_preserved(
     mock_dataset_label_unit: DataSet,
 ) -> None:
-    xr_das = mock_dataset_label_unit.to_xarray_dataarray_dict()
+    with pytest.warns(QCoDeSDeprecationWarning, match="to_xarray_dataarray_dict"):
+        xr_das = mock_dataset_label_unit.to_xarray_dataarray_dict()  # pyright: ignore[reportDeprecated]
 
     for outer_param_name, xr_da in xr_das.items():
         for param_name in xr_da.dims:
@@ -987,6 +1009,23 @@ def test_to_xarray_da_dict_paramspec_metadata_is_preserved(
             assert xr_da.attrs[spec_name] == spec_value
 
 
+def test_to_xarray_ds_dict_paramspec_metadata_is_preserved(
+    mock_dataset_label_unit: DataSet,
+) -> None:
+    xr_das = mock_dataset_label_unit.to_xarray_dataset_dict()
+
+    for outer_param_name, xr_da in xr_das.items():
+        for param_name in xr_da.dims:
+            assert xr_da.coords[param_name].attrs == _get_expected_param_spec_attrs(
+                mock_dataset_label_unit, param_name
+            )
+        expected_param_spec_attrs = _get_expected_param_spec_attrs(
+            mock_dataset_label_unit, outer_param_name
+        )
+        for spec_name, spec_value in expected_param_spec_attrs.items():
+            assert xr_da[outer_param_name].attrs[spec_name] == spec_value
+
+
 def test_export_2d_dataset(
     tmp_path_factory: TempPathFactory, mock_dataset_grid: DataSet
 ) -> None:
@@ -1597,7 +1636,7 @@ def test_geneate_pandas_index() -> None:
 @given(
     function_name=hst.sampled_from(
         [
-            "to_xarray_dataarray_dict",
+            "to_xarray_dataset_dict",
             "to_pandas_dataframe",
             "to_pandas_dataframe_dict",
             "get_parameter_data",
@@ -1642,7 +1681,7 @@ def test_export_lazy_load(
 @given(
     function_name=hst.sampled_from(
         [
-            "to_xarray_dataarray_dict",
+            "to_xarray_dataset_dict",
             "to_pandas_dataframe",
             "to_pandas_dataframe_dict",
             "get_parameter_data",