Merge pull request #2701 from AllenInstitute/ticket-2490_2

fixes issue with loading trials from BehaviorEcephysSession.from_nwb

allensdk/brain_observatory/ecephys/behavior_ecephys_session.py

@@ -1,4 +1,4 @@
-from typing import Optional, List, Dict, Any, Union, Callable, Tuple
+from typing import Optional, List, Dict, Any, Union, Callable, Tuple, Type
 
 import numpy as np
 import pandas as pd
@@ -50,6 +50,10 @@ def _get_monitor_delay():
         # used to collect the ecephys sessions for VBN
         return 0.02
 
+    @classmethod
+    def _trials_class(cls) -> Type[Trials]:
+        return VBNTrials
+
     @classmethod
     def from_lims(cls, behavior_session_id: int,
                   lims_db: Optional[Any] = None,

allensdk/test/brain_observatory/ecephys/test_behavior_ecephys_session.py

@@ -122,3 +122,17 @@ def test_getters_sanity(behavior_ecephys_session_fixture):
     behavior_ecephys_session_fixture.get_performance_metrics()
     behavior_ecephys_session_fixture.get_rolling_performance_df()
     behavior_ecephys_session_fixture.get_reward_rate()
+
+
+@pytest.mark.requires_bamboo
+def test_getters_sanity_from_nwb(
+        behavior_ecephys_session_fixture):
+    """Sanity check to make sure that the BehaviorEcephysSession
+    can use the BehaviorSession base class getter methods when read from nwb
+    """
+    nwbfile, _ = behavior_ecephys_session_fixture.to_nwb()
+
+    sess = BehaviorEcephysSession.from_nwb(nwbfile=nwbfile)
+    sess.get_performance_metrics()
+    sess.get_rolling_performance_df()
+    sess.get_reward_rate()