connect

Author: xadupre

_unittests/ut_torch_models/test_hghub_model.py

@@ -6,6 +6,7 @@
     hide_stdout,
     requires_torch,
     requires_transformers,
+    ignore_errors,
 )
 from onnx_diagnostic.torch_models.hghub.model_inputs import (
     config_class_from_architecture,
@@ -102,6 +103,7 @@ def test_get_untrained_model_with_inputs_text2text_generation(self):
     @hide_stdout()
     @requires_torch("2.7", "reduce test time")
     @requires_transformers("4.50", "reduce test time")
+    @ignore_errors(OSError)  # connectivity issues may happen
     def test_get_untrained_model_Ltesting_models(self):
         # UNHIDE=1 python _unittests/ut_torch_models/test_hghub_model.py -k L -f
         def _diff(c1, c2):

onnx_diagnostic/ext_test_case.py

@@ -108,6 +108,34 @@ def call_f(self):
     return wrapper
 
 
+def ignore_errors(errors: Union[Exception, Tuple[Exception]]) -> Callable:
+    """
+    Catches exception, skip the test if the error is expected sometimes.
+
+    :param errors: errors to ignore
+    """
+
+    def wrapper(fct):
+        if errors is None:
+            raise AssertionError(f"errors cannot be None for '{fct}'.")
+
+        def call_f(self):
+            try:
+                return fct(self)
+            except errors as e:
+                raise unittest.SkipTest(  # noqa: B904
+                    f"expecting error {e.__class__.__name__}: {e}"
+                )
+
+        try:  # noqa: SIM105
+            call_f.__name__ = fct.__name__
+        except AttributeError:
+            pass
+        return call_f
+
+    return wrapper
+
+
 def hide_stdout(f: Optional[Callable] = None) -> Callable:
     """
     Catches warnings, hides standard output.