Feature/SOF-7804-1 Update: add my org and accounts helpers

src/py/mat3ra/api_client/client.py

@@ -1,7 +1,11 @@
-from typing import Any, Optional, Tuple
+import os
+import re
+from typing import Any, List, Optional, Tuple
 
+import requests
 from pydantic import BaseModel, ConfigDict
 
+from .constants import ACCESS_TOKEN_ENV_VAR, _build_base_url
 from .endpoints.bank_materials import BankMaterialEndpoints
 from .endpoints.bank_workflows import BankWorkflowEndpoints
 from .endpoints.jobs import JobEndpoints
@@ -26,8 +30,15 @@ class APIClient(BaseModel):
     def model_post_init(self, __context: Any) -> None:
         self.my_account = Account(client=self)
         self.account = self.my_account
+        self._my_organization: Optional[Account] = None
         self._init_endpoints(self.timeout_seconds)
 
+    @property
+    def my_organization(self) -> Optional[Account]:
+        if self._my_organization is None:
+            self._my_organization = self.get_default_organization()
+        return self._my_organization
+
     @classmethod
     def env(cls) -> APIEnv:
         return APIEnv.from_env()
@@ -37,39 +48,17 @@ def auth_env(cls) -> AuthEnv:
         return AuthEnv.from_env()
 
     def _init_endpoints(self, timeout_seconds: int) -> None:
-        kwargs = {"timeout": timeout_seconds, "auth": self.auth}
-        account_id = self.auth.account_id or ""
-        auth_token = self.auth.auth_token or ""
-        self._init_core_endpoints(kwargs, account_id, auth_token)
-        self._init_bank_endpoints(kwargs, account_id, auth_token)
-
-    def _init_core_endpoints(self, kwargs: dict, account_id: str, auth_token: str) -> None:
-        self.materials = MaterialEndpoints(
-            self.host, self.port, account_id, auth_token, version=self.version, secure=self.secure, **kwargs
-        )
-        self.workflows = WorkflowEndpoints(
-            self.host, self.port, account_id, auth_token, version=self.version, secure=self.secure, **kwargs
-        )
-        self.jobs = JobEndpoints(
-            self.host, self.port, account_id, auth_token, version=self.version, secure=self.secure, **kwargs
-        )
-        self.projects = ProjectEndpoints(
-            self.host, self.port, account_id, auth_token, version=self.version, secure=self.secure, **kwargs
-        )
-        self.properties = PropertiesEndpoints(
-            self.host, self.port, account_id, auth_token, version=self.version, secure=self.secure, **kwargs
-        )
-        self.metaproperties = MetaPropertiesEndpoints(
-            self.host, self.port, account_id, auth_token, version=self.version, secure=self.secure, **kwargs
-        )
-
-    def _init_bank_endpoints(self, kwargs: dict, account_id: str, auth_token: str) -> None:
-        self.bank_materials = BankMaterialEndpoints(
-            self.host, self.port, account_id, auth_token, version=self.version, secure=self.secure, **kwargs
-        )
-        self.bank_workflows = BankWorkflowEndpoints(
-            self.host, self.port, account_id, auth_token, version=self.version, secure=self.secure, **kwargs
-        )
+        base_args = (self.host, self.port, self.auth.account_id or "", self.auth.auth_token or "")
+        base_kwargs = {"version": self.version, "secure": self.secure, "timeout": timeout_seconds, "auth": self.auth}
+
+        self.materials = MaterialEndpoints(*base_args, **base_kwargs)
+        self.workflows = WorkflowEndpoints(*base_args, **base_kwargs)
+        self.jobs = JobEndpoints(*base_args, **base_kwargs)
+        self.projects = ProjectEndpoints(*base_args, **base_kwargs)
+        self.properties = PropertiesEndpoints(*base_args, **base_kwargs)
+        self.metaproperties = MetaPropertiesEndpoints(*base_args, **base_kwargs)
+        self.bank_materials = BankMaterialEndpoints(*base_args, **base_kwargs)
+        self.bank_workflows = BankWorkflowEndpoints(*base_args, **base_kwargs)
 
     @staticmethod
     def _resolve_config(
@@ -122,9 +111,72 @@ def authenticate(
             auth_token: Optional[str] = None,
             timeout_seconds: int = 60,
     ) -> "APIClient":
-        host_value, port_value, version_value, secure_value = cls._resolve_config(host, port, version, secure,
-                                                                                  cls.env())
+        host_value, port_value, version_value, secure_value = cls._resolve_config(
+            host, port, version, secure, cls.env()
+        )
         auth = cls._auth_from_env(access_token=access_token, account_id=account_id, auth_token=auth_token)
         cls._validate_auth(auth)
-        return cls(host=host_value, port=port_value, version=version_value, secure=secure_value, auth=auth,
-                   timeout_seconds=timeout_seconds)
+        return cls(
+            host=host_value,
+            port=port_value,
+            version=version_value,
+            secure=secure_value,
+            auth=auth,
+            timeout_seconds=timeout_seconds,
+        )
+
+    def _fetch_user_data(self) -> dict:
+        access_token = self.auth.access_token or os.environ.get(ACCESS_TOKEN_ENV_VAR)
+        if not access_token:
+            raise ValueError("Access token is required to fetch user data")
+
+        url = _build_base_url(self.host, self.port, self.secure, "/api/v1/users/me")
+        response = requests.get(url, headers={"Authorization": f"Bearer {access_token}"}, timeout=30)
+        response.raise_for_status()
+        return response.json()["data"]["user"]
+
+    def _fetch_user_accounts(self) -> List[dict]:
+        return self._fetch_user_data().get("accounts", [])
+
+    def list_accounts(self) -> List[dict]:
+        accounts = self._fetch_user_accounts()
+        return [
+            {
+                "id": acc["entity"]["_id"],
+                "name": acc["entity"].get("name", ""),
+                "type": acc["entity"].get("type", "user"),
+                "isDefault": acc.get("isDefault", False),
+            }
+            for acc in accounts
+        ]
+
+    def get_account(self, name: Optional[str] = None, index: Optional[int] = None) -> Account:
+        """Get account by name (partial regex match) or index from the list of user accounts."""
+        if name is None and index is None:
+            raise ValueError("Either 'name' or 'index' must be provided")
+
+        accounts = self._fetch_user_accounts()
+
+        if index is not None:
+            return Account(client=self, id_cache=accounts[index]["entity"]["_id"])
+
+        pattern = re.compile(name, re.IGNORECASE)
+        matches = [acc for acc in accounts if pattern.search(acc["entity"].get("name", ""))]
+
+        if not matches:
+            raise ValueError(f"No account found matching '{name}'")
+        if len(matches) > 1:
+            names = [acc["entity"].get("name", "") for acc in matches]
+            raise ValueError(f"Multiple accounts match '{name}': {names}")
+
+        return Account(client=self, id_cache=matches[0]["entity"]["_id"])
+
+    def get_default_organization(self) -> Optional[Account]:
+        accounts = self._fetch_user_accounts()
+        organizations = [acc for acc in accounts if acc["entity"].get("type") == "organization"]
+
+        if not organizations:
+            return None
+
+        default_org = next((org for org in organizations if org.get("isDefault")), organizations[0])
+        return Account(client=self, id_cache=default_org["entity"]["_id"])

src/py/mat3ra/api_client/endpoints/jobs.py

@@ -54,7 +54,8 @@ def terminate(self, id_):
         """
         self.request("POST", "/".join((self.name, id_, "submit")), headers=self.headers)
 
-    def build_config(self, material_ids, workflow_id, project_id, owner_id, name, compute=None, is_multi_material=False):
+    def build_config(self, material_ids, workflow_id, project_id, owner_id, name, compute=None,
+                     is_multi_material=False):
         """
         Returns a job config based on the given parameters.
 

src/py/mat3ra/api_client/models.py

@@ -1,10 +1,9 @@
 import os
 from typing import Any, Optional
 
-import requests
 from pydantic import BaseModel, ConfigDict, Field
 
-from .constants import ACCESS_TOKEN_ENV_VAR, ACCOUNT_ID_ENV_VAR, AUTH_TOKEN_ENV_VAR, _build_base_url
+from .constants import ACCESS_TOKEN_ENV_VAR, ACCOUNT_ID_ENV_VAR, AUTH_TOKEN_ENV_VAR
 
 
 class AuthContext(BaseModel):
@@ -53,14 +52,11 @@ def _resolve_account_id(self) -> str:
             self.client.auth.account_id = account_id
             return account_id
 
-        access_token = self.client.auth.access_token or os.environ.get(ACCESS_TOKEN_ENV_VAR)
-        if not access_token:
+        if not (self.client.auth.access_token or os.environ.get(ACCESS_TOKEN_ENV_VAR)):
             raise ValueError("ACCOUNT_ID is not set and no OIDC access token is available.")
 
-        url = _build_base_url(self.client.host, self.client.port, self.client.secure, "/api/v1/users/me")
-        response = requests.get(url, headers={"Authorization": f"Bearer {access_token}"}, timeout=30)
-        response.raise_for_status()
-        account_id = response.json()["data"]["user"]["entity"]["defaultAccountId"]
+        user_data = self.client._fetch_user_data()
+        account_id = user_data["entity"]["defaultAccountId"]
         os.environ[ACCOUNT_ID_ENV_VAR] = account_id
         self.client.auth.account_id = account_id
         return account_id

tests/py/unit/test_client.py

@@ -17,6 +17,28 @@
 ME_ACCOUNT_ID = "my-account-id"
 USERS_ME_RESPONSE = {"data": {"user": {"entity": {"defaultAccountId": ME_ACCOUNT_ID}}}}
 
+ACCOUNTS_RESPONSE = {
+    "data": {
+        "user": {
+            "entity": {"defaultAccountId": ME_ACCOUNT_ID},
+            "accounts": [
+                {
+                    "entity": {"_id": "user-acc-1", "name": "John Doe", "type": "user"},
+                    "isDefault": True,
+                },
+                {
+                    "entity": {"_id": "org-acc-1", "name": "Acme Corp", "type": "organization"},
+                    "isDefault": True,
+                },
+                {
+                    "entity": {"_id": "org-acc-2", "name": "Beta Industries", "type": "organization"},
+                    "isDefault": False,
+                },
+            ],
+        }
+    }
+}
+
 
 class APIClientUnitTest(EndpointBaseUnitTest):
     def _base_env(self):
@@ -27,9 +49,9 @@ def _base_env(self):
             "API_SECURE": API_SECURE_FALSE,
         }
 
-    def _mock_users_me(self, mock_get):
+    def _mock_users_me(self, mock_get, response=None):
         mock_resp = mock.Mock()
-        mock_resp.json.return_value = USERS_ME_RESPONSE
+        mock_resp.json.return_value = response or USERS_ME_RESPONSE
         mock_resp.raise_for_status.return_value = None
         mock_get.return_value = mock_resp
 
@@ -71,3 +93,40 @@ def test_my_account_id_fetches_and_caches(self, mock_get):
             self.assertEqual(mock_get.call_args[1]["headers"]["Authorization"], f"Bearer {OIDC_ACCESS_TOKEN}")
             self.assertEqual(mock_get.call_args[1]["timeout"], 30)
             self.assertEqual(os.environ.get("ACCOUNT_ID"), ME_ACCOUNT_ID)
+
+    @mock.patch("requests.get")
+    def test_list_accounts(self, mock_get):
+        env = self._base_env() | {"OIDC_ACCESS_TOKEN": OIDC_ACCESS_TOKEN}
+        with mock.patch.dict("os.environ", env, clear=True):
+            self._mock_users_me(mock_get, ACCOUNTS_RESPONSE)
+            client = APIClient.authenticate()
+            accounts = client.list_accounts()
+
+            self.assertEqual(len(accounts), 3)
+            self.assertEqual(accounts[0]["id"], "user-acc-1")
+            self.assertEqual(accounts[0]["name"], "John Doe")
+            self.assertEqual(accounts[0]["type"], "user")
+            self.assertTrue(accounts[0]["isDefault"])
+            self.assertEqual(accounts[1]["id"], "org-acc-1")
+            self.assertEqual(accounts[1]["name"], "Acme Corp")
+            self.assertEqual(accounts[1]["type"], "organization")
+
+    @mock.patch("requests.get")
+    def test_get_account(self, mock_get):
+        env = self._base_env() | {"OIDC_ACCESS_TOKEN": OIDC_ACCESS_TOKEN}
+        with mock.patch.dict("os.environ", env, clear=True):
+            self._mock_users_me(mock_get, ACCOUNTS_RESPONSE)
+            client = APIClient.authenticate()
+
+            self.assertEqual(client.get_account(index=1).id_cache, "org-acc-1")
+            self.assertEqual(client.get_account(name="Acme").id_cache, "org-acc-1")
+            self.assertEqual(client.get_account(name="Beta.*").id_cache, "org-acc-2")
+
+    @mock.patch("requests.get")
+    def test_my_organization(self, mock_get):
+        env = self._base_env() | {"OIDC_ACCESS_TOKEN": OIDC_ACCESS_TOKEN}
+        with mock.patch.dict("os.environ", env, clear=True):
+            self._mock_users_me(mock_get, ACCOUNTS_RESPONSE)
+            client = APIClient.authenticate()
+            org = client.my_organization
+            self.assertEqual(org.id_cache, "org-acc-1")