Handle additional connection errors in preparation tests

Signed-off-by: elronbandel <[email protected]>
IBM · Feb 18, 2025 · ab96ee2 · ab96ee2
1 parent 05647c5
commit ab96ee2
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/catalog/test_preparation.py b/tests/catalog/test_preparation.py
@@ -4,7 +4,7 @@
 import tracemalloc
 
 import psutil
-from huggingface_hub.utils import GatedRepoError
+from huggingface_hub.errors import GatedRepoError, HfHubHTTPError
 from unitxt.loaders import MissingKaggleCredentialsError
 from unitxt.logging_utils import get_logger
 from unitxt.settings_utils import get_constants, get_settings
@@ -81,8 +81,8 @@ def test_preparations(self):
                         current_exception = e
                         import requests
                         while current_exception:
-                            if isinstance(current_exception, requests.exceptions.ReadTimeout):
-                                logger.error(f"ReadTimeout occurred in {file}")
+                            if isinstance(current_exception, (requests.exceptions.ReadTimeout, HfHubHTTPError)):
+                                logger.error(f"Connection error occurred in {file}. Error: {e}.")
                                 break
                             current_exception = current_exception.__cause__ or current_exception.__context__
                         else: