Write tests for alignment.py

src/beelabel/alignment.py

@@ -12,9 +12,12 @@
 from scipy.spatial.transform import Rotation as R
 import retrodetect
 import hashlib
+from typing import Union
 
+TimeOfDay = Union[str, time.struct_time]
 
-def grab_photos_in_timerange(path, starttime, endtime) -> list[str]:
+
+def grab_photos_in_timerange(path: str, starttime: TimeOfDay, endtime: TimeOfDay) -> list[str]:
     """
     Finds all numpy files (with timestamps in) at 'path' location. Avoid
     trailing / in this parameter.
@@ -27,9 +30,9 @@ def grab_photos_in_timerange(path, starttime, endtime) -> list[str]:
     Example:
         grab_photos_in_timerange('photos/system001','08:00:00','08:10:00')
     """
-    if type(starttime) == str:
+    if isinstance(starttime, str):
         starttime = time.strptime(starttime, '%H:%M:%S')
-    if type(endtime) == str:
+    if isinstance(endtime, str):
         endtime = time.strptime(endtime, '%H:%M:%S')
     output = list()
 

tests/test_alignment.py

@@ -1,13 +1,53 @@
+"""
+Unit tests for the code in beelabel/alignment.py
+"""
+
+import datetime
+import logging
+import time
+import tempfile
+from pathlib import Path
+
+import pytest
+
 import beelabel.alignment
 
 
-def test_grab_photos_in_timerange():
-    # TODO generate empty files in a temp dir
+class TestPhotosInTimeRange:
+    """
+    Tests for the grab_photos_in_timerange() function
+    """
+
+    @staticmethod
+    @pytest.fixture(scope='class', autouse=True)
+    def photo_dir():
+        with tempfile.TemporaryDirectory() as temp_dir:
+            timestamp = datetime.datetime.now()
+
+            for i in range(3):
+                timestamp + datetime.timedelta(hours=i)
+                # TODO revert to colons '%H:%M:%S'
+                time_of_day = timestamp.time().strftime('%H-%M-%S')
+                filename = f"file_{i}_{time_of_day}.np"
+                path = Path(temp_dir).joinpath(filename)
+                path.touch(exist_ok=True)
+
+            yield path
+
+    def test_grab_photos_in_timerange_str(self, photo_dir):
+        # TODO generate empty files in a temp dir
+
+        beelabel.alignment.grab_photos_in_timerange(
+            path=str(photo_dir),
+            starttime='08:00:00',
+            endtime='08:10:00'
+        )
 
-    beelabel.alignment.grab_photos_in_timerange(
-        path='photos/system001',
-        starttime='08:00:00',
-        endtime='08:10:00'
-    )
+        # TODO delete temp files
 
-    # TODO delete temp files
+    def test_grab_photos_in_timerange_struct_time(self, photo_dir):
+        beelabel.alignment.grab_photos_in_timerange(
+            path=str(photo_dir),
+            starttime=time.strptime('08:00:00', '%H:%M:%S'),
+            endtime=time.strptime('08:10:00', '%H:%M:%S')
+        )