From ad6ff06184a702a32fa28901203bfd0199df0c12 Mon Sep 17 00:00:00 2001
From: Richard Nemeth <ryshoooo@gmail.com>
Date: Sun, 15 Dec 2024 20:10:27 +0100
Subject: [PATCH] fix: retry upon 401

---
 src/keycloak/openid_connection.py | 32 +++++++++++++++++++++++++++++++
 tests/test_keycloak_admin.py      | 10 ++++------
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/src/keycloak/openid_connection.py b/src/keycloak/openid_connection.py
index 2a2e077..bf601bb 100644
--- a/src/keycloak/openid_connection.py
+++ b/src/keycloak/openid_connection.py
@@ -389,6 +389,10 @@ class KeycloakOpenIDConnection(ConnectionManager):
         """
         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 @@ class KeycloakOpenIDConnection(ConnectionManager):
         """
         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 @@ class KeycloakOpenIDConnection(ConnectionManager):
         """
         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 @@ class KeycloakOpenIDConnection(ConnectionManager):
         """
         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 @@ class KeycloakOpenIDConnection(ConnectionManager):
         """
         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 @@ class KeycloakOpenIDConnection(ConnectionManager):
         """
         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 @@ class KeycloakOpenIDConnection(ConnectionManager):
         """
         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 @@ class KeycloakOpenIDConnection(ConnectionManager):
         """
         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
diff --git a/tests/test_keycloak_admin.py b/tests/test_keycloak_admin.py
index 04179a8..11fc2f3 100644
--- a/tests/test_keycloak_admin.py
+++ b/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"):