diff --git a/atlasapi/atlas.py b/atlasapi/atlas.py index 9b9f95f..ee14c8c 100644 --- a/atlasapi/atlas.py +++ b/atlasapi/atlas.py @@ -17,7 +17,6 @@ Core module which provides access to MongoDB Atlas Cloud Provider APIs """ -import errors from atlasapi.network import Network from atlasapi.errors import * @@ -1890,7 +1889,7 @@ class _Organizations: def __init__(self, atlas): self.atlas = atlas - def _get_orgs(self, org_name: str = None, org_id: str = None) -> tuple: + def _get_orgs(self, org_name: str = None, org_id: str = None) -> Iterator[dict]: """Helper method for returning Org info from the API Args: @@ -1913,13 +1912,12 @@ def _get_orgs(self, org_name: str = None, org_id: str = None) -> tuple: raise e if not org_id: - result_list = response.get("results", []) - total_count = response.get("TotalCount", 0) + for page in response: + for each in page.get("results"): + yield each else: - result_list = [response] - total_count = 1 - - return result_list, total_count + for each in response: + yield each @property def organizations(self) -> Iterable[Organization]: @@ -1930,8 +1928,7 @@ def organizations(self) -> Iterable[Organization]: Returns (Iterable[Organization]): Yields Organization Objects. """ - result_list = self._get_orgs()[0] - for each in result_list: + for each in self._get_orgs(): yield Organization.from_dict(each) def organization_by_name(self, org_name: str) -> Organization: @@ -1944,11 +1941,10 @@ def organization_by_name(self, org_name: str) -> Organization: Returns (Organization): a single Organization object. """ - result = self._get_orgs(org_name=org_name)[0] - if isinstance(result, list): - return Organization.from_dict(result[0]) - else: - return Organization.from_dict(result) + return_list = [] + for each in self._get_orgs(org_name=org_name): + return_list.append(each) + return Organization.from_dict(return_list[0]) def organization_by_id(self, org_id: str) -> Organization: """Single organization searched by org_id. @@ -1958,12 +1954,12 @@ def organization_by_id(self, org_id: str) -> Organization: Returns (Organization): a single Organization object. """ - result = self._get_orgs(org_id=org_id) - if isinstance(result[0], list): - logger.error("returning a list?") - return Organization.from_dict(result[0][0]) - else: - return Organization.from_dict(result[0]) + + return_list = [] + for each in self._get_orgs(org_id=org_id): + return_list.append(each) + + return Organization.from_dict(return_list[0]) @property def count(self) -> int: @@ -1972,7 +1968,10 @@ def count(self) -> int: Returns (int): """ - return self._get_orgs()[1] + count = 0 + for each in self.organizations: + count += 1 + return count def get_all_projects_for_org(self, org_id: str) -> Iterable[Project]: """Get projects related to the current Org @@ -1991,8 +1990,9 @@ def get_all_projects_for_org(self, org_id: str) -> Iterable[Project]: ORG_ID=org_id) response = self.atlas.network.get(Settings.BASE_URL + uri) - for entry in response.get('results', []): - yield Project.from_dict(entry) + for page in response: + for each_project in page.get("results"): + yield Project.from_dict(each_project) class AtlasPagination: @@ -2095,26 +2095,3 @@ class CloudBackupSnapshotsGetAll(AtlasPagination): def __init__(self, atlas, pageNum, itemsPerPage): super().__init__(atlas, atlas.CloudBackups.get_all_backup_snapshots, pageNum, itemsPerPage) - - -class OrganizationProjectsGetAll(AtlasPagination): - """Pagination for Database User : Get All""" - - def __init__(self, atlas: Atlas, org_id: str, pageNum, itemsPerPage): - super().__init__(atlas, self.fetch, pageNum, itemsPerPage) - print(f"Got to init") - self._get_all_projects_for_org = atlas.Organizations.get_all_projects_for_org - self.org_id = org_id - - def fetch(self, pageNum, itemsPerPage): - """Intermediate fetching - - Args: - pageNum (int): Page number - itemsPerPage (int): Number of Events per Page - - Returns: - dict: Response payload - """ - print("Got to fetch") - return self._get_all_projects_for_org(self.org_id, pageNum, itemsPerPage) diff --git a/atlasapi/clusters.py b/atlasapi/clusters.py index d9f1f96..aa90b81 100644 --- a/atlasapi/clusters.py +++ b/atlasapi/clusters.py @@ -10,7 +10,7 @@ from typing import List, Optional from datetime import datetime import logging -from lib import DefaultReadConcerns +from atlasapi.lib import DefaultReadConcerns import pytz import uuid diff --git a/atlasapi/organizations.py b/atlasapi/organizations.py index 9d5d5fd..85d0812 100644 --- a/atlasapi/organizations.py +++ b/atlasapi/organizations.py @@ -3,6 +3,14 @@ class Organization: def __init__(self, name: str, is_deleted: bool = False, links: Optional[list] = None, id: Optional[str] = None): + """A single atlas organization, the top organizational level with Atlas + + Args: + name (str): + is_deleted (bool): + links (List[str]): + id (str): + """ self.id = id self.links = links self.is_deleted = is_deleted @@ -10,5 +18,13 @@ def __init__(self, name: str, is_deleted: bool = False, links: Optional[list] = @classmethod def from_dict(cls, data_dict: dict): + """Creates class from an Atlas returned dict. + + Args: + data_dict (dict): As returned by the atlas api. + + Returns: + + """ return cls(data_dict.get("name"), data_dict.get("isDeleted", False), data_dict.get("links", []), data_dict.get("id", None)) diff --git a/tests/test_organizations.py b/tests/test_organizations.py index 8f3c597..948e956 100644 --- a/tests/test_organizations.py +++ b/tests/test_organizations.py @@ -22,6 +22,7 @@ class ProjectTests(BaseTests): def test_00_get_organizations(self): for each in self.a.Organizations.organizations: + pprint(each.__dict__) self.assertIsInstance(each, Organization, "An Atlas should be returned") test_00_get_organizations.basic = True @@ -31,7 +32,7 @@ def test_01_get_organization_by_name(self): org_name = each.name result = self.a.Organizations.organization_by_name(org_name=org_name) - #pprint(result.__dict__) + print(f"✅Found organization {org_name}") self.assertIsInstance(result, Organization, "An Atlas should be returned") self.assertEqual(org_name, result.name, "Returned result was not the same.") @@ -42,14 +43,15 @@ def test_02_get_organization_by_id(self): org_id = each.id result = self.a.Organizations.organization_by_id(org_id) - #pprint(result.__dict__) self.assertIsInstance(result, Organization, "An Atlas should be returned") self.assertEqual(org_id, result.id, "Returned result was not the same.") + print(f"✅Found a {type(result)} with id {result.id}, when searching bu {org_id} ({result.name})") test_02_get_organization_by_id.basic = True def test_03_get_organization_count(self): result = self.a.Organizations.count + print(f"✅THere are {result} organizations.") self.assertIsInstance(result, int, "The count should be an int") test_03_get_organization_count.basic = True @@ -58,14 +60,12 @@ def test_04_get_all_projects_for_org(self): org_id = '5ac52173d383ad0caf52e11c' project_count = 0 for each_project in self.a_owner.Organizations.get_all_projects_for_org(org_id=org_id): - print(f"Found Project :{each_project.name}, {type(each_project)}") - self.assertIsInstance(each_project, Project, f"The return type was not , it was {type(each_project)}") - project_count +=1 + print(f"✅Found Project :{each_project.name}, {type(each_project)}") + self.assertIsInstance(each_project, Project, + f"The return type was not , it was {type(each_project)}") + project_count += 1 - self.assertGreater(project_count,0, "Did not find any projects, this is a bug, or the test org is not set up " - "correctly.") + self.assertGreater(project_count, 0, "Did not find any projects, this is a bug, or the test org is not set up " + "correctly.") test_04_get_all_projects_for_org.basic = True - - -