fix: retry upon 401

Author: ryshoooo

src/keycloak/openid_connection.py

@@ -389,6 +389,10 @@ def raw_get(self, *args, **kwargs):
         """
         self._refresh_if_required()
         r = super().raw_get(*args, **kwargs)
+        if r.status_code == 401:
+            self.refresh_token()
+            r = super().raw_get(*args, **kwargs)
+
         return r
 
     def raw_post(self, *args, **kwargs):
@@ -406,6 +410,10 @@ def raw_post(self, *args, **kwargs):
         """
         self._refresh_if_required()
         r = super().raw_post(*args, **kwargs)
+        if r.status_code == 401:
+            self.refresh_token()
+            r = super().raw_post(*args, **kwargs)
+
         return r
 
     def raw_put(self, *args, **kwargs):
@@ -423,6 +431,10 @@ def raw_put(self, *args, **kwargs):
         """
         self._refresh_if_required()
         r = super().raw_put(*args, **kwargs)
+        if r.status_code == 401:
+            self.refresh_token()
+            r = super().raw_put(*args, **kwargs)
+
         return r
 
     def raw_delete(self, *args, **kwargs):
@@ -440,6 +452,10 @@ def raw_delete(self, *args, **kwargs):
         """
         self._refresh_if_required()
         r = super().raw_delete(*args, **kwargs)
+        if r.status_code == 401:
+            self.refresh_token()
+            r = super().raw_delete(*args, **kwargs)
+
         return r
 
     async def a_get_token(self):
@@ -496,6 +512,10 @@ async def a_raw_get(self, *args, **kwargs):
         """
         await self.a__refresh_if_required()
         r = await super().a_raw_get(*args, **kwargs)
+        if r.status_code == 401:
+            await self.a_refresh_token()
+            r = await super().a_raw_get(*args, **kwargs)
+
         return r
 
     async def a_raw_post(self, *args, **kwargs):
@@ -513,6 +533,10 @@ async def a_raw_post(self, *args, **kwargs):
         """
         await self.a__refresh_if_required()
         r = await super().a_raw_post(*args, **kwargs)
+        if r.status_code == 401:
+            await self.a_refresh_token()
+            r = await super().a_raw_post(*args, **kwargs)
+
         return r
 
     async def a_raw_put(self, *args, **kwargs):
@@ -530,6 +554,10 @@ async def a_raw_put(self, *args, **kwargs):
         """
         await self.a__refresh_if_required()
         r = await super().a_raw_put(*args, **kwargs)
+        if r.status_code == 401:
+            await self.a_refresh_token()
+            r = await super().a_raw_put(*args, **kwargs)
+
         return r
 
     async def a_raw_delete(self, *args, **kwargs):
@@ -547,4 +575,8 @@ async def a_raw_delete(self, *args, **kwargs):
         """
         await self.a__refresh_if_required()
         r = await super().a_raw_delete(*args, **kwargs)
+        if r.status_code == 401:
+            await self.a_refresh_token()
+            r = await super().a_raw_delete(*args, **kwargs)
+
         return r

tests/test_keycloak_admin.py

@@ -2771,9 +2771,8 @@ def test_auto_refresh(admin_frozen: KeycloakAdmin, realm: str):
         "Content-Type": "application/json",
     }
 
-    with pytest.raises(KeycloakAuthenticationError) as err:
-        admin.get_realm(realm_name=realm)
-    assert err.match('401: b\'{"error":"HTTP 401 Unauthorized".*}\'')
+    res = admin.get_realm(realm_name=realm)
+    assert res["realm"] == realm
 
     # Freeze time to simulate the access token expiring
     with freezegun.freeze_time("2023-02-25 10:05:00"):
@@ -6008,9 +6007,8 @@ async def test_a_auto_refresh(admin_frozen: KeycloakAdmin, realm: str):
         "Content-Type": "application/json",
     }
 
-    with pytest.raises(KeycloakAuthenticationError) as err:
-        await admin.a_get_realm(realm_name=realm)
-    assert err.match('401: b\'{"error":"HTTP 401 Unauthorized".*}\'')
+    res = await admin.a_get_realm(realm_name=realm)
+    assert res["realm"] == realm
 
     # Freeze time to simulate the access token expiring
     with freezegun.freeze_time("2023-02-25 10:05:00"):