Debug session retrieval and unittests.

AllenInstitute · Aug 15, 2023 · badd5ea · badd5ea
1 parent 3633ce0
commit badd5ea
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 24 deletions.
diff --git a/...ior/behavior_project_cache/project_apis/data_io/behavior_neuropixels_project_cloud_api.py b/...ior/behavior_project_cache/project_apis/data_io/behavior_neuropixels_project_cloud_api.py
@@ -52,10 +52,16 @@ def get_behavior_session(
         -------
         BehaviorSession
         """
-        row = self._behavior_session_table.query(
-            f"behavior_session_id=={behavior_session_id}"
-        )
-        if row.shape[0] != 1:
+        try:
+            row = self._behavior_session_table.loc[behavior_session_id]
+        except KeyError:
+            raise RuntimeError(
+                "The behavior_session_table should have "
+                "1 and only 1 entry for a given "
+                "behavior_session_id. No sessions found for "
+                f"id={behavior_session_id}"
+            )
+        if len(row.shape) != 1:
             raise RuntimeError(
                 "The behavior_session_table should have "
                 "1 and only 1 entry for a given "
@@ -71,18 +77,7 @@ def get_behavior_session(
         if row[self.cache.file_id_column] < 0 or np.isnan(
             row[self.cache.file_id_column]
         ):
-            row = self._ecephys_session_table.query(
-                f"index=={ecephys_session_id}"
-            )
-
-        if len(row) == 0:
-            raise RuntimeError(
-                f"behavior_session: {behavior_session_id} "
-                f"corresponding to "
-                f"ecephys_session: {ecephys_session_id}"
-                f"does not exist in the behavior_session  "
-                "or ecephys_session tables."
-            )
+            row = self._ecephys_session_table.loc[ecephys_session_id]
 
         file_id = str(int(row[self.cache.file_id_column]))
         data_path = self._get_data_path(file_id=file_id)
@@ -105,14 +100,20 @@ def get_ecephys_session(
         BehaviorEcephysSession
 
         """
-        session_meta = self._ecephys_session_table.query(
-            f"index=={ecephys_session_id}"
-        )
+        try:
+            session_meta = self._ecephys_session_table.loc[ecephys_session_id]
+        except KeyError:
+            raise RuntimeError(
+                "The behavior_ecephys_session_table should "
+                "have 1 and only 1 entry for a given "
+                "ecephys_session_id. Not entries found for requested "
+                f"id={ecephys_session_id}"
+            )
         probes_meta = self._probe_table[
             (self._probe_table["ecephys_session_id"] == ecephys_session_id)
             & (self._probe_table["has_lfp_data"])
         ]
-        if session_meta.shape[0] != 1:
+        if len(session_meta.shape) != 1:
             raise RuntimeError(
                 "The behavior_ecephys_session_table should "
                 "have 1 and only 1 entry for a given "

diff --git a/...bservatory/behavior/behavior_project_cache/test_behavior_neuropixels_project_cloud_api.py b/...bservatory/behavior/behavior_project_cache/test_behavior_neuropixels_project_cloud_api.py
@@ -79,10 +79,14 @@ def mock_cache(tmpdir):
     c = {
             "behavior_sessions": pd.DataFrame({
                 "behavior_session_id": [1, 2, 3, 4],
-                "ecephys_session_id": [10, 11, 12, 13],
+                "ecephys_session_id": pd.Series(
+                    [10, 11, 12, 13],
+                    dtype='Int64'),
                 "mouse_id": [4, 4, 2, 1]}),
             "ecephys_sessions": pd.DataFrame({
-                "ecephys_session_id": [10, 11, 12, 13],
+                "ecephys_session_id": pd.Series(
+                    [10, 11, 12, 13],
+                    dtype='Int64'),
                 "behavior_session_id": [1, 2, 3, 4],
                 "file_id": [10, 11, 12, 13]}),
             "probes": pd.DataFrame({
@@ -96,12 +100,14 @@ def mock_cache(tmpdir):
             "channels": pd.DataFrame({
                 "ecephys_channel_id": [14, 15, 16],
                 "ecephys_probe_id": [4, 4, 4],
-                "ecephys_session_id": [10, 10, 10]}),
+                "ecephys_session_id": [10, 10, 10]
+            }),
             "units": pd.DataFrame({
                 "unit_id": [204, 205, 206],
                 "ecephys_channel_id": [14, 15, 16],
                 "ecephys_probe_id": [4, 4, 4],
-                "ecephys_session_id": [10, 10, 10]}),
+                "ecephys_session_id": [10, 10, 10]
+            }),
     }
 
     # round-trip the tables through csv to pick up