HTTP Error 504: Gateway Time-out

Lightning-AI · Jan 31, 2025 · d2a755d · d2a755d
1 parent db47dfb
commit d2a755d
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tests/unittests/_helpers/wrappers.py b/tests/unittests/_helpers/wrappers.py
@@ -1,6 +1,7 @@
 import os
 from functools import wraps
 from typing import Any, Callable, Optional
+from urllib.error import HTTPError
 
 import pytest
 
@@ -47,6 +48,10 @@ def run_test(*args: Any, **kwargs: Any) -> Optional[Any]:
 
             try:
                 return function(*args, **kwargs)
+            except HTTPError as ex:
+                if ex.code != 504:  # HTTP Error 504: Gateway Time-out
+                    raise ex
+                pytest.skip(reason)
             except URLError as ex:
                 if "Error 403: Forbidden" not in str(ex) or not ALLOW_SKIP_IF_BAD_CONNECTION:
                     raise ex