fix(authent): works on multi instances setup (#224)

ForestAdmin · Jun 5, 2024 · c31db91 · c31db91
1 parent f344eaa
commit c31db91
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 8 deletions.
diff --git a/src/agent_toolkit/forestadmin/agent_toolkit/utils/authentication.py b/src/agent_toolkit/forestadmin/agent_toolkit/utils/authentication.py
@@ -21,16 +21,17 @@ def register(self, url: str, registration_token: Optional[str] = None, **kwargs)
         """
         needed to avoid bj64 encoding in authorization header
         """
-        req = self.create_registration_request(**kwargs)  # type: ignore
+        req = self.create_registration_request(**kwargs)
         if self.events:
             self.events.store("Protocol request", req)
         headers = {"content-type": "application/json"}
         if registration_token is not None:
             headers["Authorization"] = "Bearer " + registration_token
 
-        rsp = self.http_request(url, "POST", data=req.to_json(), headers=headers)  # type: ignore
+        rsp = self.http_request(url, "POST", data=req.to_json(), headers=headers)
 
-        return self.handle_registration_info(rsp)  # type: ignore
+        self.handle_registration_info(rsp)
+        self.redirect_uris = self.registration_response["redirect_uris"]
 
     def get_authorization_url(self, state: str) -> str:
         args: Dict[str, Any] = {

diff --git a/src/agent_toolkit/tests/utils/test_authentication.py b/src/agent_toolkit/tests/utils/test_authentication.py
@@ -12,23 +12,31 @@ class TestCustomClientOic(TestCase):
     def test_register_should_call_http_url_with_env_secret(self):
         mock_response = Mock()
         client = CustomClientOic()
-        client.events = Mock()
+        client.events = Mock()  # type:ignore
         body_data = json.dumps(
             {"application_type": "web", "response_types": ["code"], "grant_types": ["authorization_code"]}
         )
+
+        def handle_registration_info_mock(response):
+            client.registration_response = {  # type:ignore
+                "redirect_uris": ["https://my_project.com/forest/authentication/callback"]
+            }
+
         with patch.object(client, "http_request", return_value=mock_response) as mocked_http_request:
-            with patch.object(
-                client, "handle_registration_info", side_effect=lambda rsp: rsp
+            with patch(
+                "forestadmin.agent_toolkit.utils.authentication.CustomClientOic.handle_registration_info",
+                wraps=handle_registration_info_mock,
             ) as mocked_handle_registration:
-                response = client.register("https://api.development.forestadmin.com/oidc/reg", "env_secret")
+                client.register("https://api.development.forestadmin.com/oidc/reg", "env_secret")
 
                 mocked_http_request.assert_called_once_with(
                     "https://api.development.forestadmin.com/oidc/reg",
                     "POST",
                     data=body_data,
                     headers={"content-type": "application/json", "Authorization": "Bearer env_secret"},
                 )
-                mocked_handle_registration.assert_called_once_with(response)
+                mocked_handle_registration.assert_called_once_with(mock_response)
+        self.assertEqual(client.redirect_uris, ["https://my_project.com/forest/authentication/callback"])
 
     def test_get_authorization_url_should_return_the_correct_url(self):
         client = CustomClientOic()