From 84666d8326eb0e1b9c8992b3b879f15f37081127 Mon Sep 17 00:00:00 2001
From: Sung Yun <107272191+sungwy@users.noreply.github.com>
Date: Fri, 28 Feb 2025 03:07:08 +0000
Subject: [PATCH] status code

---
 pyiceberg/catalog/rest.py  | 2 +-
 tests/catalog/test_rest.py | 7 ++++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/pyiceberg/catalog/rest.py b/pyiceberg/catalog/rest.py
index 01d7c23480..ae00454000 100644
--- a/pyiceberg/catalog/rest.py
+++ b/pyiceberg/catalog/rest.py
@@ -148,7 +148,7 @@ def _retry_hook(retry_state: RetryCallState) -> None:
 
 
 _RETRY_ARGS = {
-    "retry": retry_if_exception_type(AuthorizationExpiredError),
+    "retry": retry_if_exception_type((AuthorizationExpiredError, UnauthorizedError)),
     "stop": stop_after_attempt(2),
     "before_sleep": _retry_hook,
     "reraise": True,
diff --git a/tests/catalog/test_rest.py b/tests/catalog/test_rest.py
index 4ad825f5b7..f2fc6ceb6b 100644
--- a/tests/catalog/test_rest.py
+++ b/tests/catalog/test_rest.py
@@ -558,7 +558,8 @@ def test_list_namespace_with_parent_200(rest_mock: Mocker) -> None:
 @pytest.mark.filterwarnings(
     "ignore:Deprecated in 0.8.0, will be removed in 1.0.0. Iceberg REST client is missing the OAuth2 server URI:DeprecationWarning"
 )
-def test_list_namespaces_token_expired(rest_mock: Mocker) -> None:
+@pytest.mark.parametrize("status_code", [401, 419])
+def test_list_namespaces_token_expired_success_on_retries(rest_mock: Mocker, status_code: int) -> None:
     new_token = "new_jwt_token"
     new_header = dict(TEST_HEADERS)
     new_header["Authorization"] = f"Bearer {new_token}"
@@ -568,12 +569,12 @@ def test_list_namespaces_token_expired(rest_mock: Mocker) -> None:
         f"{TEST_URI}v1/namespaces",
         [
             {
-                "status_code": 419,
+                "status_code": status_code,
                 "json": {
                     "error": {
                         "message": "Authorization expired.",
                         "type": "AuthorizationExpiredError",
-                        "code": 419,
+                        "code": status_code,
                     }
                 },
                 "headers": TEST_HEADERS,