Skip to content

Commit

Permalink
add test to ensure that normal path is taken by default
Browse files Browse the repository at this point in the history
  • Loading branch information
jenshnielsen committed Oct 4, 2023
1 parent 2e4b979 commit 9f43f65
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 4 deletions.
7 changes: 6 additions & 1 deletion qcodes/dataset/data_set.py
Original file line number Diff line number Diff line change
Expand Up @@ -1462,8 +1462,8 @@ def _export_as_netcdf(self, path: Path, file_name: str) -> Path:
"""Export data as netcdf to a given path with file prefix"""
import xarray as xr

file_path = path / file_name
if self._estimate_ds_size() > self._export_limit:
file_path = path / file_name
log.info(
"Dataset is expected to be larger that threshold. Using distributed export.",
extra={
Expand Down Expand Up @@ -1517,6 +1517,11 @@ def _export_as_netcdf(self, path: Path, file_name: str) -> Path:
finally:
data.close()
else:
log.info(
"Writing netcdf file directly.",
extra={"file_name": file_path},
)

file_path = super()._export_as_netcdf(path=path, file_name=file_name)
return file_path

Expand Down
2 changes: 1 addition & 1 deletion qcodes/dataset/exporters/export_to_xarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -256,7 +256,7 @@ def xarray_to_h5netcdf_with_complex_numbers(
if not compute and maybe_write_job is not None:
with TqdmCallback(desc="Combining files"):
_LOG.info(
"Writing netcdf file using Dask delayed writer",
"Writing netcdf file using Dask delayed writer.",
extra={"file_name": file_path},
)
maybe_write_job.compute()
17 changes: 15 additions & 2 deletions qcodes/tests/dataset/test_dataset_export.py
Original file line number Diff line number Diff line change
Expand Up @@ -714,7 +714,20 @@ def test_export_2d_dataset(
assert xr_ds.identical(xr_ds_reimported)


def test_export_numpy_dataset(
def test_export_dataset_small_no_delated(
tmp_path_factory: TempPathFactory, mock_dataset_numpy: DataSet, caplog
) -> None:
"""
Test that a 'small' dataset does not use the delayed export.
"""
tmp_path = tmp_path_factory.mktemp("export_netcdf")
with caplog.at_level(logging.INFO):
mock_dataset_numpy.export(export_type="netcdf", path=tmp_path, prefix="qcodes_")

assert "Writing netcdf file directly" in caplog.records[0].msg


def test_export_dataset_delayed(
tmp_path_factory: TempPathFactory, mock_dataset_numpy: DataSet, caplog
) -> None:
tmp_path = tmp_path_factory.mktemp("export_netcdf")
Expand Down Expand Up @@ -747,7 +760,7 @@ def test_export_numpy_dataset(
_assert_xarray_metadata_is_as_expected(loaded_ds, mock_dataset_numpy)


def test_export_numpy_dataset_complex(
def test_export_dataset_delayed_complex(
tmp_path_factory: TempPathFactory, mock_dataset_numpy_complex: DataSet, caplog
) -> None:
tmp_path = tmp_path_factory.mktemp("export_netcdf")
Expand Down

0 comments on commit 9f43f65

Please sign in to comment.