From 69e5d9405b4a21aef2d9a80711c88fad00babff1 Mon Sep 17 00:00:00 2001
From: Jeremy Magland <jmagland@flatironinstitute.org>
Date: Tue, 28 Nov 2023 15:25:26 -0500
Subject: [PATCH] recording extractor interface tests

---
 .../baserecordingextractorinterface.py        | 11 +++++++
 .../tools/testing/data_interface_mixins.py    | 32 +++++++++++++++----
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/src/neuroconv/datainterfaces/ecephys/baserecordingextractorinterface.py b/src/neuroconv/datainterfaces/ecephys/baserecordingextractorinterface.py
index bfb4e9991..c9f8ac2e5 100644
--- a/src/neuroconv/datainterfaces/ecephys/baserecordingextractorinterface.py
+++ b/src/neuroconv/datainterfaces/ecephys/baserecordingextractorinterface.py
@@ -248,6 +248,17 @@ def set_probe(self, probe, group_mode: Literal["by_shank", "by_probe"]):
         # But neuroconv allows "group_name" property to override spike interface "group" value
         self.recording_extractor.set_property("group_name", self.recording_extractor.get_property("group").astype(str))
 
+    def has_probe(self) -> bool:
+        """
+        Check if the recording extractor has probe information.
+
+        Returns
+        -------
+        has_probe : bool
+            True if the recording extractor has probe information.
+        """
+        return self.recording_extractor.has_probe()
+
     def align_by_interpolation(
         self,
         unaligned_timestamps: np.ndarray,
diff --git a/src/neuroconv/tools/testing/data_interface_mixins.py b/src/neuroconv/tools/testing/data_interface_mixins.py
index 09018c8cc..793e78b97 100644
--- a/src/neuroconv/tools/testing/data_interface_mixins.py
+++ b/src/neuroconv/tools/testing/data_interface_mixins.py
@@ -111,13 +111,6 @@ def test_conversion_as_lone_interface(self):
                 self.case = num
                 self.test_kwargs = kwargs
                 self.interface = self.data_interface_cls(**self.test_kwargs)
-                do_set_probe = isinstance(self.interface, BaseRecordingExtractorInterface)
-                if do_set_probe:
-                    assert isinstance(self.interface, BaseRecordingExtractorInterface)
-                    self.interface.set_probe(
-                        generate_mock_probe(num_channels=self.interface.recording_extractor.get_num_channels()),
-                        group_mode="by_shank",
-                    )
                 self.check_metadata_schema_valid()
                 self.check_conversion_options_schema_valid()
                 self.check_metadata()
@@ -476,6 +469,31 @@ def test_interface_alignment(self):
 
                 self.check_nwbfile_temporal_alignment()
 
+    def test_conversion_as_lone_interface(self):
+        interface_kwargs = self.interface_kwargs
+        if isinstance(interface_kwargs, dict):
+            interface_kwargs = [interface_kwargs]
+        for num, kwargs in enumerate(interface_kwargs):
+            with self.subTest(str(num)):
+                self.case = num
+                self.test_kwargs = kwargs
+                self.interface = self.data_interface_cls(**self.test_kwargs)
+                assert isinstance(self.interface, BaseRecordingExtractorInterface)
+                if not self.interface.has_probe():
+                    self.interface.set_probe(
+                        generate_mock_probe(num_channels=self.interface.recording_extractor.get_num_channels()),
+                        group_mode="by_shank",
+                    )
+                self.check_metadata_schema_valid()
+                self.check_conversion_options_schema_valid()
+                self.check_metadata()
+                self.nwbfile_path = str(self.save_directory / f"{self.__class__.__name__}_{num}.nwb")
+                self.run_conversion(nwbfile_path=self.nwbfile_path)
+                self.check_read_nwb(nwbfile_path=self.nwbfile_path)
+
+                # Any extra custom checks to run
+                self.run_custom_checks()
+
 
 class SortingExtractorInterfaceTestMixin(DataInterfaceTestMixin, TemporalAlignmentMixin):
     data_interface_cls: BaseSortingExtractorInterface