Allow read timeout error in preparation tests (#1615)

* Allow read timeout error in preparation tests Signed-off-by: elronbandel <[email protected]> * Handle additional connection errors in preparation tests Signed-off-by: elronbandel <[email protected]> --------- Signed-off-by: elronbandel <[email protected]>
IBM · Feb 18, 2025 · 1ba974c · 1ba974c
1 parent 334b15a
commit 1ba974c
Showing 1 changed file with 13 additions and 2 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
@@ -76,7 +76,18 @@ def test_preparations(self):
                             )
                             continue
                         self.assertTrue(False)
-                        raise
+                        raise e
+                    except Exception as e:
+                        current_exception = e
+                        import requests
+                        while current_exception:
+                            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:
+                            raise e
+
                     logger.info(f"Testing preparation file: {file} passed")
                     self.assertTrue(True)