Fix tests.

luxonis · Aug 30, 2024 · 1bb0d37 · 1bb0d37
1 parent 3e4ca49
commit 1bb0d37
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/tests/unittests/test_creators/test_detections.py b/tests/unittests/test_creators/test_detections.py
@@ -5,8 +5,8 @@
 import pytest
 
 from depthai_nodes.ml.messages import (
-    ImgDetectionsWithKeypoints,
-    ImgDetectionWithKeypoints,
+    ImgDetectionExtended,
+    ImgDetectionsExtended,
 )
 from depthai_nodes.ml.messages.creators.detection import create_detection_message
 
@@ -172,10 +172,9 @@ def test_bboxes_scores_keypoints():
 
     message = create_detection_message(bboxes, scores, None, keypoints)
 
-    assert isinstance(message, ImgDetectionsWithKeypoints)
+    assert isinstance(message, ImgDetectionsExtended)
     assert all(
-        isinstance(detection, ImgDetectionWithKeypoints)
-        for detection in message.detections
+        isinstance(detection, ImgDetectionExtended) for detection in message.detections
     )
 
     for i, detection in enumerate(message.detections):