From b9e38ecc76d5a01768f3b250a5ab7aa45750c36d Mon Sep 17 00:00:00 2001
From: Bert Blommers <bblommers@users.noreply.github.com>
Date: Sun, 7 Nov 2021 15:02:48 -0100
Subject: [PATCH] CognitoIDP:list_users() - Implement AttributesToGet-parameter
 (#4539)

---
 moto/cognitoidp/models.py                | 11 ++++++++--
 moto/cognitoidp/responses.py             |  8 ++++++-
 tests/test_cognitoidp/test_cognitoidp.py | 27 ++++++++++++++++++++++++
 3 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/moto/cognitoidp/models.py b/moto/cognitoidp/models.py
index df71a0193..5cf3ffb2f 100644
--- a/moto/cognitoidp/models.py
+++ b/moto/cognitoidp/models.py
@@ -646,7 +646,9 @@ class CognitoIdpUser(BaseModel):
         }
 
     # list_users brings back "Attributes" while admin_get_user brings back "UserAttributes".
-    def to_json(self, extended=False, attributes_key="Attributes"):
+    def to_json(
+        self, extended=False, attributes_key="Attributes", attributes_to_get=None
+    ):
         user_mfa_setting_list = []
         if self.software_token_mfa_enabled:
             user_mfa_setting_list.append("SOFTWARE_TOKEN_MFA")
@@ -654,10 +656,15 @@ class CognitoIdpUser(BaseModel):
             user_mfa_setting_list.append("SMS_MFA")
         user_json = self._base_json()
         if extended:
+            attrs = [
+                attr
+                for attr in self.attributes
+                if not attributes_to_get or attr["Name"] in attributes_to_get
+            ]
             user_json.update(
                 {
                     "Enabled": self.enabled,
-                    attributes_key: self.attributes,
+                    attributes_key: attrs,
                     "MFAOptions": [],
                     "UserMFASettingList": user_mfa_setting_list,
                 }
diff --git a/moto/cognitoidp/responses.py b/moto/cognitoidp/responses.py
index 846ebc0ef..980358631 100644
--- a/moto/cognitoidp/responses.py
+++ b/moto/cognitoidp/responses.py
@@ -341,6 +341,7 @@ class CognitoIdpResponse(BaseResponse):
         limit = self._get_param("Limit")
         token = self._get_param("PaginationToken")
         filt = self._get_param("Filter")
+        attributes_to_get = self._get_param("AttributesToGet")
         users, token = cognitoidp_backends[self.region].list_users(
             user_pool_id, limit=limit, pagination_token=token
         )
@@ -385,7 +386,12 @@ class CognitoIdpResponse(BaseResponse):
                     and compare(inherent_attributes[name](user), value)
                 )
             ]
-        response = {"Users": [user.to_json(extended=True) for user in users]}
+        response = {
+            "Users": [
+                user.to_json(extended=True, attributes_to_get=attributes_to_get)
+                for user in users
+            ]
+        }
         if token:
             response["PaginationToken"] = str(token)
         return json.dumps(response)
diff --git a/tests/test_cognitoidp/test_cognitoidp.py b/tests/test_cognitoidp/test_cognitoidp.py
index 79bce2c58..3c8e3d153 100644
--- a/tests/test_cognitoidp/test_cognitoidp.py
+++ b/tests/test_cognitoidp/test_cognitoidp.py
@@ -2210,6 +2210,33 @@ def test_list_users_when_limit_more_than_total_items():
     result.shouldnt.have.key("PaginationToken")
 
 
+@mock_cognitoidp
+def test_list_users_with_attributes_to_get():
+    conn = boto3.client("cognito-idp", "us-west-2")
+    user_pool_id = conn.create_user_pool(PoolName=str(uuid.uuid4()))["UserPool"]["Id"]
+
+    for _ in range(5):
+        conn.admin_create_user(
+            UserPoolId=user_pool_id,
+            Username=str(uuid.uuid4()),
+            UserAttributes=[
+                {"Name": "family_name", "Value": "Doe"},
+                {"Name": "given_name", "Value": "Jane"},
+                {"Name": "custom:foo", "Value": "bar"},
+            ],
+        )
+
+    result = conn.list_users(
+        UserPoolId=user_pool_id, AttributesToGet=["given_name", "custom:foo", "unknown"]
+    )
+    users = result["Users"]
+    users.should.have.length_of(5)
+    for user in users:
+        user["Attributes"].should.have.length_of(2)
+        user["Attributes"].should.contain({"Name": "given_name", "Value": "Jane"})
+        user["Attributes"].should.contain({"Name": "custom:foo", "Value": "bar"})
+
+
 @mock_cognitoidp
 def test_admin_disable_user():
     conn = boto3.client("cognito-idp", "us-west-2")