Added support to the other tests and fixed the remaining tests

Author: vaibhavatlan

pyatlan/client/atlan.py

@@ -867,7 +867,7 @@ def get_group_by_name(
         self,
         alias: str,
         limit: int = 20,
-    ) -> Optional[List[AtlanGroup]]:
+    ) -> Optional[GroupResponse]:
         """Deprecated - use group.get_by_name() instead."""
         warn(
             "This method is deprecated, please use 'group.get_by_name' instead, which offers identical functionality.",
@@ -1002,7 +1002,7 @@ def get_users_by_email(
         self,
         email: str,
         limit: int = 20,
-    ) -> Optional[List[AtlanUser]]:
+    ) -> Optional[UserResponse]:
         """Deprecated - use user.get_by_email() instead."""
         warn(
             "This method is deprecated, please use 'user.get_by_email' instead, which offers identical functionality.",

pyatlan/client/group.py

@@ -155,15 +155,15 @@ def get_all(
         response: GroupResponse = self.get(
             offset=offset, limit=limit, sort=sort, columns=columns
         )
-        return response  # type: ignore
+        return response
 
     @validate_arguments
     def get_by_name(
         self,
         alias: str,
         limit: int = 20,
         offset: int = 0,
-    ) -> Optional[List[AtlanGroup]]:
+    ) -> Optional[GroupResponse]:
         """
         Retrieve all groups with a name that contains the provided string.
         (This could include a complete group name, in which case there should be at most
@@ -175,13 +175,12 @@ def get_by_name(
         :param offset: starting point for the list of groups when paging
         :returns: all groups whose name (in the UI) contains the provided string
         """
-        if response := self.get(
+        response: GroupResponse = self.get(
             offset=offset,
             limit=limit,
             post_filter='{"$and":[{"alias":{"$ilike":"%' + alias + '%"}}]}',
-        ):
-            return response.records
-        return None
+        )
+        return response
 
     @validate_arguments
     def get_members(

pyatlan/client/user.py

@@ -48,7 +48,7 @@ def __init__(self, client: ApiCaller):
     @validate_arguments
     def create(
         self, users: List[AtlanUser], return_info: bool = False
-    ) -> Optional[List[AtlanUser]]:
+    ) -> Optional[UserResponse]:
         """
         Create one or more new users.
 
@@ -219,7 +219,7 @@ def get_by_email(
         email: str,
         limit: int = 20,
         offset: int = 0,
-    ) -> Optional[List[AtlanUser]]:
+    ) -> Optional[UserResponse]:
         """
         Retrieves all users with email addresses that contain the provided email.
         (This could include a complete email address, in which case there should be at
@@ -232,21 +232,20 @@ def get_by_email(
         :param offset: starting point for the list of users when pagin
         :returns: all users whose email addresses contain the provided string
         """
-        if response := self.get(
+        response: UserResponse = self.get(
             offset=offset,
             limit=limit,
             post_filter='{"email":{"$ilike":"%' + email + '%"}}',
-        ):
-            return response.records
-        return None
+        )
+        return response
 
     @validate_arguments
     def get_by_emails(
         self,
         emails: List[str],
         limit: int = 20,
         offset: int = 0,
-    ) -> Optional[List[AtlanUser]]:
+    ) -> Optional[UserResponse]:
         """
         Retrieves all users with email addresses that match the provided list of emails.
 
@@ -256,9 +255,10 @@ def get_by_emails(
         :returns: list of users whose email addresses match the provided list
         """
         email_filter = '{"email":{"$in":' + dumps(emails or [""]) + "}}"
-        if response := self.get(offset=offset, limit=limit, post_filter=email_filter):
-            return response.records
-        return None
+        response: UserResponse = self.get(
+            offset=offset, limit=limit, post_filter=email_filter
+        )
+        return response
 
     @validate_arguments
     def get_by_username(self, username: str) -> Optional[AtlanUser]:
@@ -281,7 +281,7 @@ def get_by_username(self, username: str) -> Optional[AtlanUser]:
     @validate_arguments
     def get_by_usernames(
         self, usernames: List[str], limit: int = 5, offset: int = 0
-    ) -> Optional[List[AtlanUser]]:
+    ) -> Optional[UserResponse]:
         """
         Retrieves users based on their usernames.
 
@@ -291,11 +291,10 @@ def get_by_usernames(
         :returns: the users with the specified usernames
         """
         username_filter = '{"username":{"$in":' + dumps(usernames or [""]) + "}}"
-        if response := self.get(
+        response: UserResponse = self.get(
             offset=offset, limit=limit, post_filter=username_filter
-        ):
-            return response.records
-        return None
+        )
+        return response.records  # type: ignore
 
     @validate_arguments
     def add_to_groups(

pyatlan/samples/search/and_star_assets.py

@@ -46,8 +46,8 @@ def list_users_in_group(name: str) -> List[str]:
     """
     usernames: List[str] = []
     if groups := client.group.get_by_name(alias=name):
-        if groups[0].id is not None and (
-            response := client.group.get_members(guid=groups[0].id)
+        if groups.records[0].id is not None and (  # type: ignore
+            response := client.group.get_members(guid=groups.records[0].id)  # type: ignore
         ):
             if response.records and len(response.records) > 0:
                 usernames.extend(

tests/integration/admin_test.py

@@ -52,8 +52,9 @@ def test_create_group(client: AtlanClient, group: CreateGroupResponse):
     assert group
     r = client.group.get_by_name(GROUP_NAME)
     assert r
-    assert len(r) == 1
-    group1_full = r[0]
+    assert r.records is not None
+    assert len(r.records) == 1
+    group1_full = r.records[0]
     assert group1_full
     assert group1_full.path
     assert group1_full.name
@@ -102,8 +103,9 @@ def test_group_get_pagination(client: AtlanClient, group: CreateGroupResponse):
 def test_group_get_members_pagination(client: AtlanClient, group: CreateGroupResponse):
     groups = client.group.get_by_name(alias=GROUP_NAME)
     assert groups
-    assert len(groups) == 1
-    group1 = groups[0]
+    assert groups.records
+    assert len(groups.records) == 1
+    group1 = groups.records[0]
     assert group1.id
     response = client.group.get_members(guid=group1.id, request=UserRequest(limit=1))
 
@@ -177,29 +179,33 @@ def test_retrieve_existing_user(client: AtlanClient, group: CreateGroupResponse)
     assert user1.group_count == 1 + _default_group_count
     response = client.user.get_by_usernames(usernames=[FIXED_USER])
     assert response
-    assert len(response) == 1
-    fixed_user = response[0]
+    assert response.records is not None
+    assert len(response.records) == 1
+    fixed_user = response.records[0]
     assert fixed_user
     assert fixed_user.id
     users_list = client.user.get_by_usernames(usernames=[])
     assert users_list == []
     users_list = client.user.get_by_email(EMAIL_DOMAIN)
     assert users_list
-    assert len(users_list) >= 1
+    assert users_list.records is not None
+    assert len(users_list.records) >= 1
     email = user1.email
     assert email
     users_list = client.user.get_by_email(email)
     assert users_list
-    assert len(users_list) == 1
-    assert user1.email == users_list[0].email
-    assert user1.username == users_list[0].username
-    assert user1.attributes == users_list[0].attributes
+    assert users_list.records is not None
+    assert len(users_list.records) == 1
+    assert user1.email == users_list.records[0].email
+    assert user1.username == users_list.records[0].username
+    assert user1.attributes == users_list.records[0].attributes
     users_list = client.user.get_by_emails(emails=[email])
     assert users_list
-    assert len(users_list) == 1
-    assert user1.email == users_list[0].email
-    assert user1.username == users_list[0].username
-    assert user1.attributes == users_list[0].attributes
+    assert users_list.records is not None
+    assert len(users_list.records) == 1
+    assert user1.email == users_list.records[0].email
+    assert user1.username == users_list.records[0].username
+    assert user1.attributes == users_list.records[0].attributes
     users_list = client.user.get_by_emails(emails=[])
     assert users_list == []
 
@@ -208,8 +214,9 @@ def test_retrieve_existing_user(client: AtlanClient, group: CreateGroupResponse)
 def test_update_groups(client: AtlanClient, group: CreateGroupResponse):
     groups = client.group.get_by_name(alias=GROUP_NAME)
     assert groups
-    assert len(groups) == 1
-    group1 = groups[0]
+    assert groups.records is not None
+    assert len(groups.records) == 1
+    group1 = groups.records[0]
     group1.attributes = AtlanGroup.Attributes(description=["Now with a description!"])
     client.group.update(group1)
 
@@ -221,8 +228,9 @@ def test_updated_groups(
 ):
     groups = client.group.get_by_name(alias=GROUP_NAME)
     assert groups
-    assert len(groups) == 1
-    group1 = groups[0]
+    assert groups.records is not None
+    assert len(groups.records) == 1
+    group1 = groups.records[0]
     assert group1
     assert group1.id == group.group
     assert group1.attributes
@@ -237,8 +245,9 @@ def test_remove_user_from_group(
 ):
     groups = client.group.get_by_name(alias=GROUP_NAME)
     assert groups
-    assert len(groups) == 1
-    group1 = groups[0]
+    assert groups.records is not None
+    assert len(groups.records) == 1
+    group1 = groups.records[0]
     assert group1.id
     fixed_user = client.user.get_by_username(FIXED_USER)
     assert fixed_user

tests/integration/custom_metadata_test.py

@@ -500,12 +500,14 @@ def _get_groups(
 ) -> Tuple[AtlanGroup, AtlanGroup]:
     candidates = client.group.get_by_name(GROUP_NAME1)
     assert candidates
-    assert len(candidates) == 1
-    group1 = candidates[0]
+    assert candidates.records is not None
+    assert len(candidates.records) == 1
+    group1 = candidates.records[0]
     candidates = client.group.get_by_name(GROUP_NAME2)
     assert candidates
-    assert len(candidates) == 1
-    group2 = candidates[0]
+    assert candidates.records is not None
+    assert len(candidates.records) == 1
+    group2 = candidates.records[0]
     return group1, group2
 
 

tests/integration/suggestions_test.py

@@ -438,10 +438,12 @@ def owner_group(
     response = client.group.create(group=to_create)
     assert response
     group = client.group.get_by_name(PREFIX)
-    assert group and len(group) == 1
-    assert group[0].id
-    yield group[0]
-    client.group.purge(group[0].id)
+    assert group
+    assert group.records is not None
+    assert len(group.records) == 1
+    assert group.records[0].id
+    yield group.records[0]
+    client.group.purge(group.records[0].id)
 
 
 def test_connection(client: AtlanClient, connection: Connection):

tests/integration/test_sso_client.py

@@ -40,8 +40,9 @@ def group(client: AtlanClient) -> Generator[AtlanGroup, None, None]:
     g = client.group.create(group=to_create, user_ids=[str(fixed_user.id)])
     groups = client.group.get_by_name(alias=GROUP_NAME)
     assert groups
-    assert len(groups) == 1
-    yield groups[0]
+    assert groups.records is not None
+    assert len(groups.records) == 1
+    yield groups.records[0]
     assert g.group
     delete_group(client, g.group)
 

tests/unit/test_client.py

@@ -1816,8 +1816,8 @@ def test_user_create_with_info(
     test_users = [AtlanUser.create(email="[email protected]", role_name="$member")]
     response = client.create(users=test_users, return_info=True)
 
-    assert len(response) == 1
-    user = response[0]
+    assert len(response.current_page()) == 1
+    user = response.current_page()[0]
     assert user
     assert user.username
     assert user.email