feat: improve test coverage

AllenNeuralDynamics · Dec 1, 2023 · 947f712 · 947f712
1 parent c3a9c8d
commit 947f712
Show file tree

Hide file tree

Showing 26 changed files with 759 additions and 970 deletions.
diff --git a/examples/bergamo_ophys_session.py b/examples/bergamo_ophys_session.py
@@ -2,8 +2,7 @@
 
 import datetime
 
-from aind_data_schema.core.session import FieldOfView, Session, Stream
-from aind_data_schema.models.device_configurations import DetectorConfigs, LaserConfigs
+from aind_data_schema.core.session import DetectorConfigs, FieldOfView, LaserConfigs, Session, Stream
 from aind_data_schema.models.modalities import Modality
 from aind_data_schema.models.stimulus import PhotoStimulation, PhotoStimulationGroup, StimulusEpoch
 

diff --git a/examples/ephys_session.py b/examples/ephys_session.py
@@ -2,9 +2,16 @@
 
 from datetime import datetime
 
-from aind_data_schema.core.session import DomeModule, EphysModule, ManipulatorModule, Session, Stream
+from aind_data_schema.core.session import (
+    DomeModule,
+    EphysModule,
+    EphysProbeConfigs,
+    LaserConfigs,
+    ManipulatorModule,
+    Session,
+    Stream,
+)
 from aind_data_schema.models.coordinates import Coordinates3d
-from aind_data_schema.models.device_configurations import EphysProbeConfigs, LaserConfigs
 from aind_data_schema.models.modalities import Modality
 
 red_laser = LaserConfigs(name="Red Laser", wavelength=700, excitation_power=100)

diff --git a/examples/ophys_session.py b/examples/ophys_session.py
@@ -2,8 +2,7 @@
 
 import datetime
 
-from aind_data_schema.core.session import Session, Stream
-from aind_data_schema.models.device_configurations import DetectorConfigs, FiberConnectionConfigs, LaserConfigs
+from aind_data_schema.core.session import DetectorConfigs, FiberConnectionConfigs, LaserConfigs, Session, Stream
 from aind_data_schema.models.modalities import Modality
 
 t = datetime.datetime(2022, 7, 12, 7, 00, 00)

diff --git a/src/aind_data_schema/base.py b/src/aind_data_schema/base.py
@@ -2,9 +2,9 @@
 
 import re
 from pathlib import Path
-from typing import Final, Optional
+from typing import Optional
 
-from pydantic import BaseModel, ConfigDict, Extra, Field, PrivateAttr
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
 
 
 class AindModel(BaseModel):
@@ -26,11 +26,6 @@ class AindCoreModel(AindModel):
         ..., pattern=r"^\d+.\d+.\d+$", description="schema version", title="Version", frozen=True
     )
 
-    # @classmethod
-    # def default_file_extension(cls) -> str:
-    #     """Public method to retrieve protected _DEFAULT_FILE_EXTENSION"""
-    #     return cls._DEFAULT_FILE_EXTENSION.default
-
     @classmethod
     def default_filename(cls):
         """
@@ -55,7 +50,7 @@ def write_standard_file(self, output_directory: Optional[Path] = None, prefix=No
         if prefix:
             filename = str(prefix) + "_" + filename
         if suffix:
-            filename = filename.replace(self._DEFAULT_FILE_EXTENSION, suffix)
+            filename = filename.replace(self._FILE_EXTENSION, suffix)
 
         if output_directory is not None:
             output_directory = Path(output_directory)

diff --git a/src/aind_data_schema/core/acquisition.py b/src/aind_data_schema/core/acquisition.py
@@ -9,7 +9,7 @@
 
 from aind_data_schema.base import AindCoreModel, AindModel
 from aind_data_schema.imaging.tile import AcquisitionTile
-from aind_data_schema.models.devices import Calibration, Maintenance, ImmersionMedium
+from aind_data_schema.models.devices import Calibration, ImmersionMedium, Maintenance
 from aind_data_schema.models.process_names import ProcessName
 from aind_data_schema.models.units import SizeUnit
 
@@ -131,8 +131,7 @@ def from_direction_code(cls, v: Union[str, List[Axis]]) -> List[Axis]:
 
             axes = []
             for i, c in enumerate(v):
-                axis = Axis(name=name_lookup[i], direction=direction_lookup[c],
-                            dimension=i)
+                axis = Axis(name=name_lookup[i], direction=direction_lookup[c], dimension=i)
                 axes.append(axis)
             return axes
         else: