Update tests

qubvel · Jul 18, 2024 · 558d9f7 · 558d9f7
1 parent 7047811
commit 558d9f7
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/tests/models/vitmatte/test_image_processing_vitmatte.py b/tests/models/vitmatte/test_image_processing_vitmatte.py
@@ -213,5 +213,6 @@ def test_image_processor_preprocess_arguments(self):
                 warnings.simplefilter("always")
                 image_processor(image, trimaps=trimap, extra_argument=True)
 
-            self.assertEqual(len(raised_warnings), 1)
-            self.assertIn("extra_argument", str(raised_warnings[0].message))
+            messages = " ".join([str(w.message) for w in raised_warnings])
+            self.assertGreaterEqual(len(raised_warnings), 1)
+            self.assertIn("extra_argument", messages)
diff --git a/tests/test_image_processing_common.py b/tests/test_image_processing_common.py
@@ -454,8 +454,9 @@ def test_image_processor_preprocess_arguments(self):
                     warnings.simplefilter("always")
                     image_processor(inputs, extra_argument=True)
 
-                self.assertEqual(len(raised_warnings), 1)
-                self.assertIn("extra_argument", str(raised_warnings[0].message))
+                messages = " ".join([str(w.message) for w in raised_warnings])
+                self.assertGreaterEqual(len(raised_warnings), 1)
+                self.assertIn("extra_argument", messages)
                 is_tested = True
 
         if not is_tested: