From 2390fa09c3960aa45e302e4c718684f95cc6c95b Mon Sep 17 00:00:00 2001
From: Romazes <romazes20@gmail.com>
Date: Tue, 1 Oct 2024 17:29:23 +0300
Subject: [PATCH] refactor: QCAuth0Authorization model

---
 lean/components/api/auth0_client.py       |  2 +-
 lean/models/api.py                        | 40 +++++++----------------
 lean/models/json_module.py                |  4 +--
 tests/components/api/test_auth0_client.py |  5 ++-
 4 files changed, 16 insertions(+), 35 deletions(-)

diff --git a/lean/components/api/auth0_client.py b/lean/components/api/auth0_client.py
index 1f87eb98..f1350c1f 100644
--- a/lean/components/api/auth0_client.py
+++ b/lean/components/api/auth0_client.py
@@ -45,7 +45,7 @@ def read(self, brokerage_id: str) -> QCAuth0Authorization:
 
             data = self._api.post("live/auth0/read", payload)
             # Store in cache
-            result = QCAuth0Authorization.from_raw(data)
+            result = QCAuth0Authorization(**data)
             self._cache[brokerage_id] = result
             return result
         except RequestFailedError as e:
diff --git a/lean/models/api.py b/lean/models/api.py
index 3572d31b..99a69509 100644
--- a/lean/models/api.py
+++ b/lean/models/api.py
@@ -22,25 +22,8 @@
 # The keys of properties are not changed, so they don't obey the rest of the project's naming conventions
 
 
-class Account(WrappedBaseModel):
-    id: str
-    name: str
-
-
-class Authorization(WrappedBaseModel):
-    client_info: Dict[str, str]  # Holds flexible client data like client ID and tokens
-    accounts: List[Account]  # Holds account information
-
-    @classmethod
-    def from_raw(cls, raw_data: Dict[str, any]) -> 'Authorization':
-        """
-        Custom constructor to preprocess raw authorization data
-        Separates client info from accounts.
-        """
-        # Separate the client info (non-list values) from accounts (list of accounts)
-        client_info = {k: v for k, v in raw_data.items() if not isinstance(v, list)}
-        accounts = raw_data.get('accounts', [])
-        return cls(client_info=client_info, accounts=accounts)
+class QCAuth0Authorization(WrappedBaseModel):
+    authorization: Optional[Dict[str, Any]]
 
     def get_account_ids(self) -> List[str]:
         """
@@ -52,20 +35,19 @@ def get_account_ids(self) -> List[str]:
         Returns:
             List[str]: A list of account IDs.
         """
-        return [account.id for account in self.accounts] if self.accounts else []
+        accounts = self.authorization.get('accounts', [])
+        return [account["id"] for account in accounts] if accounts else []
 
+    def get_authorization_config_without_account(self) -> Dict[str, str]:
+        """
+        Returns the authorization data without the 'accounts' key.
 
-class QCAuth0Authorization(WrappedBaseModel):
-    authorization: Optional[Authorization]
+        Iterates through the 'authorization' dictionary and excludes the 'accounts' entry.
 
-    @classmethod
-    def from_raw(cls, raw_data: Dict[str, any]) -> 'QCAuth0Authorization':
-        """
-        Custom constructor to preprocess raw QCAuth0Authorization data
+        Returns:
+            Dict[str, str]: Authorization details excluding 'accounts'.
         """
-        authorization_data = raw_data.get('authorization', {})
-        authorization = Authorization.from_raw(authorization_data) if authorization_data else None
-        return cls(authorization=authorization)
+        return {key: value for key, value in self.authorization.items() if key != 'accounts'}
 
 
 class ProjectEncryptionKey(WrappedBaseModel):
diff --git a/lean/models/json_module.py b/lean/models/json_module.py
index 95f5463a..fa93c1bf 100644
--- a/lean/models/json_module.py
+++ b/lean/models/json_module.py
@@ -217,11 +217,11 @@ def config_build(self,
             elif isinstance(configuration, AuthConfiguration):
                 auth_authorizations = get_authorization(container.api_client.auth0, self._display_name.lower(), logger)
                 logger.debug(f'auth: {auth_authorizations}')
-                configuration._value = auth_authorizations.authorization.client_info
+                configuration._value = auth_authorizations.get_authorization_config_without_account()
                 for inner_config in self._lean_configs:
                     if any(condition._dependent_config_id == configuration._id for condition in
                            inner_config._filter._conditions):
-                        api_account_ids = auth_authorizations.authorization.get_account_ids()
+                        api_account_ids = auth_authorizations.get_account_ids()
                         config_dash = inner_config._id.replace('-', '_')
                         inner_config._choices = api_account_ids
                         if user_provided_options and config_dash in user_provided_options:
diff --git a/tests/components/api/test_auth0_client.py b/tests/components/api/test_auth0_client.py
index 9e9838be..dcf240bf 100644
--- a/tests/components/api/test_auth0_client.py
+++ b/tests/components/api/test_auth0_client.py
@@ -44,6 +44,5 @@ def test_auth0client() -> None:
 
     assert result
     assert result.authorization
-    assert len(result.authorization.client_info) > 0
-    assert len(result.authorization.accounts) > 0
-    assert len(result.authorization.get_account_ids()) > 0
+    assert len(result.authorization) > 0
+    assert len(result.get_account_ids()) > 0