From 558d9f7789ea98275ff464d3d80b718bb68ba768 Mon Sep 17 00:00:00 2001 From: Pavel Iakubovskii Date: Thu, 18 Jul 2024 12:36:32 +0000 Subject: [PATCH] Update tests --- tests/models/vitmatte/test_image_processing_vitmatte.py | 5 +++-- tests/test_image_processing_common.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/models/vitmatte/test_image_processing_vitmatte.py b/tests/models/vitmatte/test_image_processing_vitmatte.py index 638bcfb958ae89..288ed53d190d00 100644 --- 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 index fa267cedf8d3b5..a3fbb6eeb78178 100644 --- 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: