Fix get_inspection

equinor · Oct 25, 2024 · 010fc94 · 010fc94
1 parent 31db3b9
commit 010fc94
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 22 deletions.
diff --git a/src/isar_robot/inspections.py b/src/isar_robot/inspections.py
@@ -56,7 +56,7 @@ def create_image(task_actions: Union[TakeImage, TakeThermalImage]):
     filepath: Path = random.choice(list(example_images.iterdir()))
     image.data = _read_data_from_file(filepath)
 
-    return [image]
+    return image
 
 
 def create_video(task_actions: TakeVideo):
@@ -76,7 +76,7 @@ def create_video(task_actions: TakeVideo):
     filepath: Path = random.choice(list(example_videos.iterdir()))
     video.data = _read_data_from_file(filepath)
 
-    return [video]
+    return video
 
 
 def create_thermal_video(task_actions: TakeThermalVideo):
@@ -96,7 +96,7 @@ def create_thermal_video(task_actions: TakeThermalVideo):
     filepath: Path = random.choice(list(example_thermal_videos.iterdir()))
     thermal_video.data = _read_data_from_file(filepath)
 
-    return [thermal_video]
+    return thermal_video
 
 
 def create_audio(task_actions: RecordAudio):
@@ -116,7 +116,7 @@ def create_audio(task_actions: RecordAudio):
     filepath: Path = random.choice(list(example_thermal_videos.iterdir()))
     audio.data = _read_data_from_file(filepath)
 
-    return [audio]
+    return audio
 
 
 def _read_data_from_file(filename: Path) -> bytes:

diff --git a/tests/test_inspections.py b/tests/test_inspections.py
@@ -13,44 +13,32 @@
 def test_create_image():
     task_actions = TakeImage(target=target)
 
-    list_of_images = inspections.create_image(task_actions)
+    inspection_image = inspections.create_image(task_actions)
 
-    assert len(list_of_images) == 1
-
-    inspection_image = list_of_images[0]
     assert inspection_image.metadata.file_type == "jpg"
 
 
 def test_create_video():
     task_actions = TakeImage(target=target)
 
-    list_of_videos = inspections.create_video(task_actions)
-
-    assert len(list_of_videos) == 1
+    inspection_video = inspections.create_video(task_actions)
 
-    inspection_video = list_of_videos[0]
     assert inspection_video.metadata.file_type == "mp4"
 
 
 def test_create_thermal_video():
     task_actions = TakeThermalVideo(target=target, duration=10)
 
-    list_of_thermal_videos = inspections.create_thermal_video(task_actions)
+    inspection_video = inspections.create_thermal_video(task_actions)
 
-    assert len(list_of_thermal_videos) == 1
-
-    inspection_video = list_of_thermal_videos[0]
     assert inspection_video.metadata.file_type == "mp4"
     assert inspection_video.metadata.duration == 10
 
 
 def test_create_audio():
     task_actions = RecordAudio(target=target, duration=10)
 
-    list_of_audio_recordings = inspections.create_audio(task_actions)
-
-    assert len(list_of_audio_recordings) == 1
+    inspection_recording = inspections.create_audio(task_actions)
 
-    inspection_recordings = list_of_audio_recordings[0]
-    assert inspection_recordings.metadata.file_type == "wav"
-    assert inspection_recordings.metadata.duration == 10
+    assert inspection_recording.metadata.file_type == "wav"
+    assert inspection_recording.metadata.duration == 10