fix test

SainsburyWellcomeCentre · Jul 3, 2024 · de63859 · de63859
1 parent cb4b7c0
commit de63859
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/test_unit/test_tracking_utils.py b/tests/test_unit/test_tracking_utils.py
@@ -49,8 +49,9 @@ def test_write_tracked_bbox_to_csv(csv_writer, csv_output):
     bbox = np.array([10, 20, 50, 80, 1])
     frame = np.zeros((100, 100, 3), dtype=np.uint8)
     frame_name = "frame_0001.png"
+    pred_score = 0.900
 
-    write_tracked_bbox_to_csv(bbox, frame, frame_name, csv_writer)
+    write_tracked_bbox_to_csv(bbox, frame, frame_name, csv_writer, pred_score)
 
     expected_row = (
         "frame_0001.png",
@@ -59,7 +60,7 @@ def test_write_tracked_bbox_to_csv(csv_writer, csv_output):
         1,
         0,
         '"{""name"":""rect"",""x"":10,""y"":20,""width"":40,""height"":60}"',
-        '"{""track"":""1""}"',
+        '"{""track"":""1"", ""pred_score"":""0.9""}"',
     )
     expected_row_str = ",".join(map(str, expected_row))
     assert csv_output.getvalue().strip() == expected_row_str