Repository: aeye-lab/pymovements
Pull request #201 (2023-02-27 10:26:38)

Datasets: Write tests for Dataset.load() method

codecov-commenter: # [Codecov](https://codecov.io/gh/aeye-lab/pymovements/pull/201?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab) Report
Base: **60.47**% // Head: **66.90**% // Increases project coverage by **`+6.43%`** :tada:
> Coverage data is based on head [(`10eb931`)](https://codecov.io/gh/aeye-lab/pymovements/pull/201?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab) compared to base [(`8fe9daa`)](https://codecov.io/gh/aeye-lab/pymovements/commit/8fe9daa4c697a51915d7656752aaed5cfd4d9a38?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab).
> Patch coverage: 92.85% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #201      +/-   ##
==========================================
+ Coverage   60.47%   66.90%   +6.43%     
==========================================
  Files          21       21              
  Lines         802      834      +32     
  Branches      188      197       +9     
==========================================
+ Hits          485      558      +73     
+ Misses        306      251      -55     
- Partials       11       25      +14     
```


| [Impacted Files](https://codecov.io/gh/aeye-lab/pymovements/pull/201?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab) | Coverage Δ | |
|---|---|---|
| [src/pymovements/datasets/gazebase.py](https://codecov.io/gh/aeye-lab/pymovements/pull/201?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab#diff-c3JjL3B5bW92ZW1lbnRzL2RhdGFzZXRzL2dhemViYXNlLnB5) | `94.44% <83.33%> (-5.56%)` | :arrow_down: |
| [src/pymovements/datasets/judo1000.py](https://codecov.io/gh/aeye-lab/pymovements/pull/201?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab#diff-c3JjL3B5bW92ZW1lbnRzL2RhdGFzZXRzL2p1ZG8xMDAwLnB5) | `94.44% <85.71%> (-5.56%)` | :arrow_down: |
| [src/pymovements/datasets/dataset.py](https://codecov.io/gh/aeye-lab/pymovements/pull/201?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab#diff-c3JjL3B5bW92ZW1lbnRzL2RhdGFzZXRzL2RhdGFzZXQucHk=) | `45.16% <93.33%> (+20.04%)` | :arrow_up: |
| [src/pymovements/datasets/public\_dataset.py](https://codecov.io/gh/aeye-lab/pymovements/pull/201?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab#diff-c3JjL3B5bW92ZW1lbnRzL2RhdGFzZXRzL3B1YmxpY19kYXRhc2V0LnB5) | `59.61% <100.00%> (+12.11%)` | :arrow_up: |
| [src/pymovements/utils/paths.py](https://codecov.io/gh/aeye-lab/pymovements/pull/201?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab#diff-c3JjL3B5bW92ZW1lbnRzL3V0aWxzL3BhdGhzLnB5) | `54.54% <0.00%> (+31.81%)` | :arrow_up: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/aeye-lab/pymovements/pull/201?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aeye-lab).


Patch:
diff --git a/README.md b/README.md
index 9b5f5bd..f4406c6 100644
--- a/README.md
+++ b/README.md
@@ -14,7 +14,7 @@
 [![License](https://img.shields.io/pypi/l/pymovements.svg)](https://github.com/aeye-lab/pymovements/blob/master/LICENSE.txt)
 [![Test Status](https://img.shields.io/github/actions/workflow/status/aeye-lab/pymovements/tests.yml?label=tests)](https://github.com/aeye-lab/pymovements/actions/workflows/tests.yml)
 [![Documentation Status](https://readthedocs.org/projects/pymovements/badge/?version=latest)](https://pymovements.readthedocs.io/en/latest/?badge=latest)
-[![codecov](https://codecov.io/github/aeye-lab/pymovements/branch/main/graph/badge.svg?token=QY3NDHAT2C)](https://codecov.io/github/aeye-lab/pymovements)
+[![codecov](https://codecov.io/github/aeye-lab/pymovements/branch/main/graph/badge.svg?token=QY3NDHAT2C)](https://app.codecov.io/gh/aeye-lab/pymovements)
 [![PyPI downloads/month](https://img.shields.io/pypi/dm/pymovements.svg)](https://pypistats.org/packages/pymovements)
 
 
diff --git a/src/pymovements/datasets/dataset.py b/src/pymovements/datasets/dataset.py
index 4bd93f8..2d22edf 100644
--- a/src/pymovements/datasets/dataset.py
+++ b/src/pymovements/datasets/dataset.py
@@ -99,6 +99,10 @@ class Dataset:
 
         self.experiment = experiment
 
+        if filename_regex is None:
+            raise ValueError('filename_regex must not be None')
+        if not isinstance(filename_regex, str):
+            raise TypeError('filename_regex must be of type str')
         self._filename_regex = filename_regex
 
         if filename_regex_dtypes is None:
@@ -151,12 +155,9 @@ class Dataset:
         AttributeError
             If no regular expression for parsing filenames is defined.
         RuntimeError
-            If an error occured during matching filenames or no files have been found.
+            If an error occurred during matching filenames or no files have been found.
         """
-        if self._filename_regex is not None:
-            filename_regex = re.compile(self._filename_regex)
-        else:
-            raise AttributeError('no regular expression for filenames is defined.')
+        filename_regex = re.compile(self._filename_regex)
 
         # Get all filepaths that match regular expression.
         csv_filepaths = get_filepaths(
@@ -171,19 +172,16 @@ class Dataset:
             # All csv_filepaths already match the filename_regex.
             match = filename_regex.match(filepath.name)
 
-            # This actually should never happen but mypy will complain otherwise.
+            # This should never happen but mypy will complain otherwise.
             if match is None:
                 raise RuntimeError(
-                    f'file {filepath} did not match regular expression {self._filename_regex}',
+                    f'file {filepath} did not match regular expression {filename_regex}',
                 )
 
             # We use the groupdict of the match as a base and add the filepath.
             fileinfo_dict = match.groupdict()
 
-            for fileinfo_key, fileinfo_dtype in self._filename_regex_dtypes.items():
-                fileinfo_dict[fileinfo_key] = fileinfo_dtype(fileinfo_dict[fileinfo_key])
-
-            fileinfo_dict['filepath'] = str(filepath)
+            fileinfo_dict['filepath'] = str(filepath.relative_to(self.raw_rootpath))
             fileinfo_dicts.append(fileinfo_dict)
 
         if len(fileinfo_dicts) == 0:
@@ -193,11 +191,18 @@ class Dataset:
         fileinfo_df = pl.from_dicts(dicts=fileinfo_dicts, infer_schema_length=1)
         fileinfo_df = fileinfo_df.sort(by='filepath')
 
+        fileinfo_df = fileinfo_df.with_columns([
+            pl.col(fileinfo_key).cast(fileinfo_dtype)
+            for fileinfo_key, fileinfo_dtype in self._filename_regex_dtypes.items()
+        ])
+
         return fileinfo_df
 
     def take_subset(
             self,
-            subset: None | dict[str, float | int | str | list[float | int | str]] = None,
+            subset: None | dict[
+                str, bool | float | int | str | list[bool | float | int | str],
+            ] = None,
     ) -> None:
         """Take a subset of the dataset.
 
@@ -240,8 +245,8 @@ class Dataset:
                 column_values = subset_value
             else:
                 raise TypeError(
-                    f'subset value must be of float, int, str or a list of these but key-value pair'
-                    f'{subset_key}: {subset_value} is of type {type(subset_value)}',
+                    f'subset value must be of type bool, float, int, str or a list of these but'
+                    f' key-value pair {subset_key}: {subset_value} is of type {type(subset_value)}',
                 )
 
             self.fileinfo = self.fileinfo.filter(pl.col(subset_key).is_in(column_values))
@@ -268,40 +273,42 @@ class Dataset:
         """
         self._check_fileinfo()
 
-        file_dfs: list[pl.DataFrame] = []
+        gaze_dfs: list[pl.DataFrame] = []
 
         # Read gaze files from fileinfo attribute.
         for file_id, filepath in enumerate(tqdm(self.fileinfo['filepath'])):
-            filepath = Path(filepath)
+            filepath = self.raw_rootpath / filepath
 
             if preprocessed:
                 filepath = self._raw_to_preprocessed_filepath(filepath)
 
             if filepath.suffix == '.csv':
-                file_df = pl.read_csv(filepath, **self._custom_read_kwargs)
+                gaze_df = pl.read_csv(filepath, **self._custom_read_kwargs)
             elif filepath.suffix == '.feather':
-                file_df = pl.read_ipc(filepath)
+                gaze_df = pl.read_ipc(filepath)
             else:
                 raise RuntimeError('data files of type {filepath.suffix} are not supported.')
 
             # Add fileinfo columns to dataframe.
             for column in self.fileinfo.columns[::-1]:
-                # Preprocessed data already has fileinfo columns.
-                if preprocessed:
-                    continue
-
                 if column == 'filepath':
                     continue
 
                 column_value = self.fileinfo.select(column)[file_id][0, 0]
-                file_df = file_df.select([
+                gaze_df = gaze_df.select([
                     pl.lit(column_value).alias(column),
                     pl.all(),
                 ])
 
-            file_dfs.append(file_df)
+            # Cast columns from fileinfo according to specification.
+            gaze_df = gaze_df.with_columns([
+                pl.col(fileinfo_key).cast(fileinfo_dtype)
+                for fileinfo_key, fileinfo_dtype in self._filename_regex_dtypes.items()
+            ])
 
-        return file_dfs
+            gaze_dfs.append(gaze_df)
+
+        return gaze_dfs
 
     def load_event_files(self) -> list[pl.DataFrame]:
         """Load all available event files.
@@ -322,7 +329,7 @@ class Dataset:
 
         # read and preprocess input files
         for _, filepath in enumerate(tqdm(self.fileinfo['filepath'])):
-            filepath = Path(filepath)
+            filepath = self.raw_rootpath / Path(filepath)
 
             filepath = self._raw_to_event_filepath(filepath)
 
@@ -543,11 +550,14 @@ class Dataset:
             raw_filepath = Path(self.fileinfo[file_id, 'filepath'])
             events_filepath = self._raw_to_event_filepath(raw_filepath)
 
-            events_filepath.parent.mkdir(parents=True, exist_ok=True)
+            for column in event_df.columns:
+                if column in self.fileinfo.columns:
+                    event_df.drop(column)
 
             if verbose > 2:
                 print('Save file to', events_filepath)
 
+            events_filepath.parent.mkdir(parents=True, exist_ok=True)
             event_df.write_ipc(events_filepath)
 
     def save_preprocessed(self, verbose: int = 1):
@@ -567,11 +577,14 @@ class Dataset:
             raw_filepath = Path(self.fileinfo[file_id, 'filepath'])
             preprocessed_filepath = self._raw_to_preprocessed_filepath(raw_filepath)
 
-            preprocessed_filepath.parent.mkdir(parents=True, exist_ok=True)
+            for column in gaze_df.columns:
+                if column in self.fileinfo.columns:
+                    gaze_df = gaze_df.drop(column)
 
             if verbose > 2:
                 print('Save file to', preprocessed_filepath)
 
+            preprocessed_filepath.parent.mkdir(parents=True, exist_ok=True)
             gaze_df.write_ipc(preprocessed_filepath)
 
     @property
@@ -725,7 +738,9 @@ class Dataset:
         Path
             The Path to the preprocessed feather file.
         """
-        relative_raw_dirpath = raw_filepath.parent.relative_to(self.raw_rootpath)
+        relative_raw_dirpath = raw_filepath.parent
+        relative_raw_dirpath = relative_raw_dirpath.relative_to(self.raw_rootpath)
+
         preprocessed_file_dirpath = self.preprocessed_rootpath / relative_raw_dirpath
 
         # Get new filename for saved feather file.
@@ -748,7 +763,9 @@ class Dataset:
         Path
             The Path to the event feather file.
         """
-        relative_raw_dirpath = raw_filepath.parent.relative_to(self.raw_rootpath)
+        relative_raw_dirpath = raw_filepath.parent
+        relative_raw_dirpath = relative_raw_dirpath.relative_to(self.raw_rootpath)
+
         events_file_dirpath = self.events_rootpath / relative_raw_dirpath
 
         # Get new filename for saved feather file.


Test patch:
diff --git a/tests/datasets/test_dataset.py b/tests/datasets/test_dataset.py
index 88de246..e44b3ba 100644
--- a/tests/datasets/test_dataset.py
+++ b/tests/datasets/test_dataset.py
@@ -20,6 +20,7 @@
 """Test all functionality in pymovements.datasets.dataset."""
 from pathlib import Path
 
+import numpy as np
 import polars as pl
 import pytest
 from polars.testing import assert_frame_equal
@@ -28,6 +29,235 @@ from pymovements.datasets.dataset import Dataset
 from pymovements.events.events import Event
 
 
+def create_raw_gaze_files_from_fileinfo(gaze_dfs, fileinfo, rootpath):
+    rootpath.mkdir(parents=True, exist_ok=True)
+
+    for gaze_df, fileinfo_row in zip(gaze_dfs, fileinfo.to_dicts()):
+        filepath = fileinfo_row['filepath']
+
+        for key in fileinfo_row.keys():
+            if key in gaze_df.columns:
+                gaze_df = gaze_df.drop(key)
+
+        gaze_df.write_csv(rootpath / filepath)
+
+
+def create_preprocessed_gaze_files_from_fileinfo(gaze_dfs, fileinfo, rootpath):
+    rootpath.mkdir(parents=True, exist_ok=True)
+
+    for gaze_df, fileinfo_row in zip(gaze_dfs, fileinfo.to_dicts()):
+        filepath = fileinfo_row['filepath']
+        filepath = filepath.replace('csv', 'feather')
+
+        for key in fileinfo_row.keys():
+            if key in gaze_df.columns:
+                gaze_df = gaze_df.drop(key)
+
+        gaze_df.write_ipc(rootpath / filepath)
+
+
+def create_event_files_from_fileinfo(gaze_dfs, fileinfo, rootpath):
+    rootpath.mkdir(parents=True, exist_ok=True)
+
+    for gaze_df, fileinfo_row in zip(gaze_dfs, fileinfo.to_dicts()):
+        filepath = fileinfo_row['filepath']
+        filepath = filepath.replace('csv', 'feather')
+
+        for key in fileinfo_row.keys():
+            if key in gaze_df.columns:
+                gaze_df = gaze_df.drop(key)
+
+        gaze_df.write_ipc(rootpath / filepath)
+
+
+def mock_toy(rootpath):
+    subject_ids = list(range(1, 21))
+
+    fileinfo = pl.DataFrame(data={'subject_id': subject_ids}, schema={'subject_id': pl.Int64})
+
+    fileinfo = fileinfo.with_columns([
+        pl.format('{}.csv', 'subject_id').alias('filepath'),
+    ])
+
+    fileinfo = fileinfo.sort(by='filepath')
+
+    gaze_dfs = []
+    for fileinfo_row in fileinfo.to_dicts():
+        gaze_df = pl.from_dict(
+            {
+                'subject_id': fileinfo_row['subject_id'],
+                'time': np.arange(1000),
+                'x_left_pix': np.zeros(1000),
+                'y_left_pix': np.zeros(1000),
+                'x_right_pix': np.zeros(1000),
+                'y_right_pix': np.zeros(1000),
+            },
+            schema={
+                'subject_id': pl.Int64,
+                'time': pl.Int64,
+                'x_left_pix': pl.Float64,
+                'y_left_pix': pl.Float64,
+                'x_right_pix': pl.Float64,
+                'y_right_pix': pl.Float64,
+            },
+        )
+        gaze_dfs.append(gaze_df)
+
+    create_raw_gaze_files_from_fileinfo(gaze_dfs, fileinfo, rootpath / 'raw')
+
+    preprocessed_gaze_dfs = []
+    for fileinfo_row in fileinfo.to_dicts():
+        gaze_df = pl.from_dict(
+            {
+                'subject_id': fileinfo_row['subject_id'],
+                'time': np.arange(1000),
+                'x_left_pix': np.zeros(1000),
+                'y_left_pix': np.zeros(1000),
+                'x_right_pix': np.zeros(1000),
+                'y_right_pix': np.zeros(1000),
+                'x_left_dva': np.zeros(1000),
+                'y_left_dva': np.zeros(1000),
+                'x_right_dva': np.zeros(1000),
+                'y_right_dva': np.zeros(1000),
+                'x_left_vel': np.zeros(1000),
+                'y_left_vel': np.zeros(1000),
+                'x_right_vel': np.zeros(1000),
+                'y_right_vel': np.zeros(1000),
+            },
+            schema={
+                'subject_id': pl.Int64,
+                'time': pl.Int64,
+                'x_left_pix': pl.Float64,
+                'y_left_pix': pl.Float64,
+                'x_right_pix': pl.Float64,
+                'y_right_pix': pl.Float64,
+                'x_left_dva': pl.Float64,
+                'y_left_dva': pl.Float64,
+                'x_right_dva': pl.Float64,
+                'y_right_dva': pl.Float64,
+                'x_left_vel': pl.Float64,
+                'y_left_vel': pl.Float64,
+                'x_right_vel': pl.Float64,
+                'y_right_vel': pl.Float64,
+            },
+        )
+        preprocessed_gaze_dfs.append(gaze_df)
+
+    create_preprocessed_gaze_files_from_fileinfo(
+        preprocessed_gaze_dfs, fileinfo, rootpath / 'preprocessed',
+    )
+
+    event_dfs = []
+    for fileinfo_row in fileinfo.to_dicts():
+        event_df = pl.from_dict(
+            {
+                'type': 'saccade',
+                'onset': np.arange(0, 901, 100),
+                'offset': np.arange(10, 911, 100),
+            },
+            schema={'type': pl.Utf8, 'onset': pl.Int64, 'offset': pl.Int64},
+        )
+        event_dfs.append(event_df)
+
+    create_event_files_from_fileinfo(event_dfs, fileinfo, rootpath / 'events')
+
+    return {
+        'init_kwargs': {
+            'root': rootpath,
+            'filename_regex': r'(?P<subject_id>\d+).csv',
+            'filename_regex_dtypes': {'subject_id': pl.Int64},
+        },
+        'fileinfo': fileinfo,
+        'raw_gaze_dfs': gaze_dfs,
+        'preprocessed_gaze_dfs': preprocessed_gaze_dfs,
+        'event_dfs': event_dfs,
+    }
+
+
+@pytest.fixture(name='dataset_configuration', params=['Toy'])
+def fixture_dataset(request, tmp_path):
+    rootpath = tmp_path
+
+    if request.param == 'Toy':
+        dataset_dict = mock_toy(rootpath)
+    else:
+        raise ValueError(f'{request.param} not supported as dataset mock')
+
+    yield dataset_dict
+
+
+def test_load_correct_fileinfo(dataset_configuration):
+    dataset = Dataset(**dataset_configuration['init_kwargs'])
+    dataset.load()
+
+    expected_fileinfo = dataset_configuration['fileinfo']
+    assert_frame_equal(dataset.fileinfo, expected_fileinfo)
+
+
+def test_load_correct_raw_gaze_dfs(dataset_configuration):
+    dataset = Dataset(**dataset_configuration['init_kwargs'])
+    dataset.load()
+
+    expected_gaze_dfs = dataset_configuration['raw_gaze_dfs']
+    for result_gaze_df, expected_gaze_df in zip(dataset.gaze, expected_gaze_dfs):
+        assert_frame_equal(result_gaze_df, expected_gaze_df)
+
+
+def test_load_correct_preprocessed_gaze_dfs(dataset_configuration):
+    dataset = Dataset(**dataset_configuration['init_kwargs'])
+    dataset.load(preprocessed=True)
+
+    expected_gaze_dfs = dataset_configuration['preprocessed_gaze_dfs']
+    for result_gaze_df, expected_gaze_df in zip(dataset.gaze, expected_gaze_dfs):
+        assert_frame_equal(result_gaze_df, expected_gaze_df)
+
+
+def test_load_correct_event_dfs(dataset_configuration):
+    dataset = Dataset(**dataset_configuration['init_kwargs'])
+    dataset.load(events=True)
+
+    expected_event_dfs = dataset_configuration['event_dfs']
+    for result_event_df, expected_event_df in zip(dataset.events, expected_event_dfs):
+        assert_frame_equal(result_event_df, expected_event_df)
+
+
+@pytest.mark.parametrize(
+    'init_kwargs, exception',
+    [
+        pytest.param(
+            {'root': 'data', 'filename_regex': None},
+            ValueError,
+            id='filename_regex_none_value',
+        ),
+        pytest.param(
+            {'root': 'data', 'filename_regex': 1},
+            TypeError,
+            id='filename_regex_wrong_type',
+        ),
+    ],
+)
+def test_init_exceptions(init_kwargs, exception):
+    with pytest.raises(exception):
+        Dataset(**init_kwargs)
+
+
+@pytest.mark.parametrize(
+    'init_kwargs, exception',
+    [
+        pytest.param(
+            {'root': '/not/a/real/path'},
+            RuntimeError,
+            id='no_files_present',
+        ),
+    ],
+)
+def test_load_exceptions(init_kwargs, exception):
+    dataset = Dataset(**init_kwargs)
+
+    with pytest.raises(exception):
+        dataset.load()
+
+
 @pytest.mark.parametrize(
     'events_init, events_expected',
     [
diff --git a/tests/utils/test_checks.py b/tests/utils/test_checks.py
index 08c20aa..65213cd 100644
--- a/tests/utils/test_checks.py
+++ b/tests/utils/test_checks.py
@@ -20,12 +20,132 @@
 """
 Test pymovements checks.
 """
+import numpy as np
+import pytest
 
+from pymovements.utils.checks import check_nan_both_channels
+from pymovements.utils.checks import check_no_zeros
+from pymovements.utils.checks import check_shapes_positions_velocities
 
-def test_check_no_zeros_exception():
-    # use fixtures: (single variable, list, np.array)
-    # check both valid and invalid inputs with the help of a second variable
+
+@pytest.mark.parametrize(
+    'variable, expected_error',
+    [
+        pytest.param(5, None, id='non_zero_single_variable_raises_no_error'),
+        pytest.param(0, ValueError, id='zero_single_variable_raises_value_error'),
+        pytest.param([1, 2, 3], None, id='non_zero_list_raises_no_error'),
+        pytest.param([1, 0, 3], ValueError, id='zero_list_raises_value_error'),
+        pytest.param(np.array([1, 2, 3]), None, id='non_zero_np_array_raises_no_error'),
+        pytest.param(np.array([1, 0, 3]), ValueError, id='zero_np_array_raises_value_error'),
+    ],
+)
+def test_check_no_zeros_raises_error(variable, expected_error):
+    """
+    Test that check_no_zeros() only raises an Exception if there are zeros in the input array.
+    """
+    if expected_error is None:
+        check_no_zeros(variable)
+    else:
+        with pytest.raises(expected_error):
+            check_no_zeros(variable)
+
+
+@pytest.mark.parametrize(
+    'arr, expected_error',
+    [
+        pytest.param(
+            np.array([[1, 2], [3, 4]]),
+            None,
+            id='no_nans_raises_no_error',
+        ),
+        pytest.param(
+            np.array([[1, 2], [np.nan, np.nan]]),
+            None,
+            id='nans_same_time_steps_raises_no_error',
+        ),
+        pytest.param(
+            np.array([[np.nan, 2], [np.nan, 4]]),
+            ValueError,
+            id='nans_different_time_steps_raises_value_error',
+        ),
+        pytest.param(
+            np.array([[np.nan, 2], [3, 4]]),
+            ValueError,
+            id='nans_only_left_channel_raises_value_error',
+        ),
+        pytest.param(
+            np.array([[1, np.nan], [3, 4]]),
+            ValueError,
+            id='nans_only_right_channel_raises_value_error',
+        ),
+    ],
+)
+def test_check_nan_both_channels_raises_error(arr, expected_error):
+    """
+    Test that check_nan_both_channels() only raises an Exception if all nans
+    occur at the same time step for both channels.
+    """
+    if expected_error is None:
+        check_nan_both_channels(arr)
+    else:
+        with pytest.raises(expected_error):
+            check_nan_both_channels(arr)
+
+
+# Test check_shapes_positions_velocities
+@pytest.mark.parametrize(
+    'kwargs, expected_error',
+    [
+        pytest.param(
+            {
+                'positions': np.array([[1, 2], [3, 4]]),
+                'velocities': np.array([[1, 2], [3, 4]]),
+            },
+            None,
+            id='positions_and_velocities_shape_N_2_raises_no_error',
+        ),
+        pytest.param(
+            {
+                'positions': np.array([[1, 2], [3, 4]]),
+                'velocities': np.array([1, 2, 3, 4]),
+            },
+            ValueError,
+            id='positions_shape_N_2_velocities_not_shape_N_2_raises_value_error',
+        ),
+        pytest.param(
+            {
+                'positions': np.array([1, 2, 3, 4]),
+                'velocities': np.array([[1, 2], [3, 4]]),
+            },
+            ValueError,
+            id='positions_not_shape_N_2_velocities_shape_N_2_raises_value_error',
+        ),
+        pytest.param(
+            {
+                'positions': np.array([1, 2, 3, 4]),
+                'velocities': np.array([1, 2, 3, 4]),
+            },
+            ValueError,
+            id='positions_and_velocities_not_shape_N_2_raises_value_error',
+        ),
+        pytest.param(
+            {
+                'positions': np.array([[1, 2], [3, 4]]),
+                'velocities': np.array([[1, 2], [3, 4], [5, 6]]),
+            },
+            ValueError,
+            id='positions_and_velocities_N_2_but_different_lengths_raises_value_error',
+        ),
+    ],
+)
+def test_check_shapes_positions_velocities_raises_error(kwargs, expected_error):
     """
-    Test that check_no_zeros() only raises an Exception iff there are zeros in the input array.
+    Test that check_shapes_positions_velocities() only raises an Exception if
+    the shapes of the positions and velocities are not (N, 2) or if the lengths
+    of the positions and velocities are not equal.
     """
-    return
+    if expected_error is None:
+        check_shapes_positions_velocities(**kwargs)
+    else:
+        with pytest.raises(expected_error):
+            check_shapes_positions_velocities(**kwargs)