Fix test_load_img_url_timeout (huggingface#25976)

* fix * fix * fix --------- Co-authored-by: ydshieh <[email protected]>
parambharat · Sep 26, 2023 · 6dc5f72 · 6dc5f72
1 parent 265a2a8
commit 6dc5f72
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/utils/test_image_utils.py b/tests/utils/test_image_utils.py
@@ -21,7 +21,7 @@
 import numpy as np
 import pytest
 from huggingface_hub.file_download import http_get
-from requests import ReadTimeout
+from requests import ConnectTimeout, ReadTimeout
 
 from tests.pipelines.test_pipelines_document_question_answering import INVOICE_URL
 from transformers import is_torch_available, is_vision_available
@@ -491,7 +491,7 @@ def test_load_img_url(self):
 
     @is_flaky()
     def test_load_img_url_timeout(self):
-        with self.assertRaises(ReadTimeout):
+        with self.assertRaises((ReadTimeout, ConnectTimeout)):
             load_image(INVOICE_URL, timeout=0.001)
 
     def test_load_img_local(self):