From 13077923660b8f4f27be6a00021e6bf3dfd95ec9 Mon Sep 17 00:00:00 2001 From: Daniel Ballard Date: Mon, 29 Jul 2024 11:52:56 -0400 Subject: [PATCH] Allow admin_update_profile to support a different username and preformatted list of cognito attr dicts --- pycognito/__init__.py | 36 ++++++++++++-- tests.py | 111 +++++++++++++++++++++++++++++++++++++++++- 2 files changed, 143 insertions(+), 4 deletions(-) diff --git a/pycognito/__init__.py b/pycognito/__init__.py index bd28318a..20a1ffe7 100644 --- a/pycognito/__init__.py +++ b/pycognito/__init__.py @@ -26,6 +26,21 @@ def cognito_to_dict(attr_list, attr_map=None): return attr_dict +def is_cognito_attr_list(attr_list): + """ + :param attr_list: List of User Pool attribute dicts + :return: bool indicating whether the list contains User Pool attribute dicts + """ + if not isinstance(attr_list, list): + return False + for attr_dict in attr_list: + if not isinstance(attr_dict, dict): + return False + if not attr_dict.keys() <= {"Name", "Value"}: + return False + return True + + def dict_to_cognito(attributes, attr_map=None): """ :param attributes: Dictionary of User Pool attribute names/values @@ -534,11 +549,26 @@ def logout(self): self.access_token = None self.token_type = None - def admin_update_profile(self, attrs, attr_map=None): - user_attrs = dict_to_cognito(attrs, attr_map) + def admin_update_profile(self, attrs, attr_map=None, username=None): + """ + Updates the user specified (defaults to self.username) with the provided attrs + :param attrs: Dictionary of attribute name, values + :param attr_map: Dictionary map from Cognito attributes to attribute + :param username: Username to update + :return: + """ + if username is None: + username = self.username + + # if already formatted for cognito then use as is + if not is_cognito_attr_list(attrs): + user_attrs = dict_to_cognito(attrs, attr_map) + else: + user_attrs = attrs + self.client.admin_update_user_attributes( UserPoolId=self.user_pool_id, - Username=self.username, + Username=username, UserAttributes=user_attrs, ) diff --git a/tests.py b/tests.py index 8bc0e5e0..e164c071 100644 --- a/tests.py +++ b/tests.py @@ -14,7 +14,13 @@ import requests import requests_mock -from pycognito import Cognito, UserObj, GroupObj, TokenVerificationException +from pycognito import ( + is_cognito_attr_list, + Cognito, + UserObj, + GroupObj, + TokenVerificationException, +) from pycognito.aws_srp import AWSSRP from pycognito.utils import RequestsSrpAuth @@ -496,5 +502,108 @@ def test_srp_requests_http_auth(self, m): self.assertNotEqual(access_token_orig, access_token_new) +@moto.mock_aws +class AdminUpdateProfileTestCase(unittest.TestCase): + username = "user@test.com" + password = "Testing123!" + + def setUp(self) -> None: + + cognito_idp_client = boto3.client("cognito-idp", region_name="us-east-1") + + user_pool = cognito_idp_client.create_user_pool( + PoolName="pycognito-test-pool", + AliasAttributes=[ + "email", + ], + UsernameAttributes=[ + "email", + ], + ) + self.user_pool_id = user_pool["UserPool"]["Id"] + + cognito_idp_client.admin_create_user( + UserPoolId=self.user_pool_id, + Username=self.username, + TemporaryPassword=self.password, + MessageAction="SUPPRESS", + ) + cognito_idp_client.admin_set_user_password( + UserPoolId=self.user_pool_id, + Username=self.username, + Password=self.password, + Permanent=True, + ) + client_app = cognito_idp_client.create_user_pool_client( + UserPoolId=self.user_pool_id, + ClientName="test-client", + RefreshTokenValidity=1, + AccessTokenValidity=1, + IdTokenValidity=1, + TokenValidityUnits={ + "AccessToken": "hour", + "IdToken": "hour", + "RefreshToken": "days", + }, + ) + self.client_id = client_app["UserPoolClient"]["ClientId"] + + def test_cognito_attr_list(self): + assert ( + is_cognito_attr_list( + [{"Name": "given_name", "Value": "John", "Bad_Key": "Test"}] + ) + is False + ) + assert ( + is_cognito_attr_list( + [{"Name": "given_name", "Value": "John"}, {"Bad_Key": "Test"}] + ) + is False + ) + assert ( + is_cognito_attr_list( + [ + {"Name": "given_name", "Value": "John"}, + {"Name": "middle_name", "Value": "Smith"}, + ] + ) + is True + ) + + def test_user_admin_update_profile(self): + cognito = Cognito( + user_pool_id=self.user_pool_id, + username=self.username, + client_id=self.client_id, + ) + cognito.authenticate(self.password) + + # verify attr dict from cognito_to_dict and dict_to_cognito works as expected + cognito.admin_update_profile(attrs={"given_name": "John"}) + assert cognito.get_user().given_name == "John" + + # verify attr_map works as expected + cognito.admin_update_profile( + attrs={"gn": "Steve"}, attr_map={"given_name": "gn"} + ) + assert cognito.get_user().given_name == "Steve" + + # verify that cognito formatted list works as expected + cognito.admin_update_profile( + attrs=[{"Name": "given_name", "Value": "Bob"}], username=self.username + ) + assert cognito.get_user().given_name == "Bob" + + # is_cognito_attr_list returns False so dict_to_cognito is called which will fail on this bad list input + try: + cognito.admin_update_profile( + attrs=[{"Name": "given_name", "Value": "John", "Bad_Key": "Test"}], + username=self.username, + ) + except AttributeError as err: + self.assertEqual(str(err), "'list' object has no attribute 'items'") + + if __name__ == "__main__": unittest.main()