diff --git a/mdps_ds_lib/lib/aws/aws_cognito.py b/mdps_ds_lib/lib/aws/aws_cognito.py index e873e77..0818a8b 100644 --- a/mdps_ds_lib/lib/aws/aws_cognito.py +++ b/mdps_ds_lib/lib/aws/aws_cognito.py @@ -18,3 +18,50 @@ def get_groups(self, username: str): return [] belonged_groups = [k['GroupName'] for k in response['Groups']] return belonged_groups + + def add_user_to_group(self, username: str, group_name: str): + # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/cognito-idp/client/admin_add_user_to_group.html + + response = self.__cognito.admin_add_user_to_group( + UserPoolId=self.__user_pool_id, + Username=username, + GroupName=group_name, + ) + if response['ResponseMetadata']['HTTPStatusCode'] != 200: + raise RuntimeError(response) + return response + + def remove_user_from_group(self, username: str, group_name: str): + # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/cognito-idp/client/admin_remove_user_from_group.html + response = self.__cognito.admin_remove_user_from_group( + UserPoolId=self.__user_pool_id, + Username=username, + GroupName=group_name, + ) + if response['ResponseMetadata']['HTTPStatusCode'] != 200: + raise RuntimeError(response) + return response + + def add_group(self, group_name: str): + # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/cognito-idp/client/create_group.html + response = self.__cognito.create_group( + GroupName=group_name, + UserPoolId=self.__user_pool_id, + # Description='NA', + # RoleArn='string', + ) + if response['ResponseMetadata']['HTTPStatusCode'] != 200: + raise RuntimeError(response) + return response + + def delete_group(self, group_name: str): + # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/cognito-idp/client/delete_group.html + response = self.__cognito.delete_group( + GroupName=group_name, + UserPoolId=self.__user_pool_id, + # Description='NA', + # RoleArn='string', + ) + if response['ResponseMetadata']['HTTPStatusCode'] != 200: + raise RuntimeError(response) + return response diff --git a/tests/mdps_ds_lib/lib/aws/test_aws_cognito.py b/tests/mdps_ds_lib/lib/aws/test_aws_cognito.py index 192b028..98558bf 100644 --- a/tests/mdps_ds_lib/lib/aws/test_aws_cognito.py +++ b/tests/mdps_ds_lib/lib/aws/test_aws_cognito.py @@ -5,9 +5,22 @@ class TestAwsCognitor(TestCase): def test_01(self): - cognito = AwsCognito('us-west-2_FLDyXE2mO') - wphyo_groups = cognito.get_groups('wphyo') + cognito = AwsCognito('us-west-2_yaOw3yj0z') + sample_group_name = 'UNIT_TEST_GROUP_WPHYO' + username = 'wphyo' + result = cognito.add_group(sample_group_name) + print(result) + result = cognito.add_user_to_group(username, sample_group_name) + print(result) + wphyo_groups = cognito.get_groups(username) self.assertTrue(isinstance(wphyo_groups, list), f'response is not list. {wphyo_groups}') self.assertTrue(len(wphyo_groups) > 0, f'empty list') - print(wphyo_groups) + self.assertTrue(sample_group_name in wphyo_groups, f'empty list') + result = cognito.remove_user_from_group(username, sample_group_name) + print(result) + wphyo_groups = cognito.get_groups(username) + self.assertTrue(isinstance(wphyo_groups, list), f'response is not list. {wphyo_groups}') + self.assertTrue(sample_group_name not in wphyo_groups, f'empty list') + result = cognito.delete_group(sample_group_name) + print(result) return