Implement user-group relationships for cognito-idp

moto/cognitoidp/models.py

@@ -196,6 +196,10 @@ class CognitoIdpGroup(BaseModel):
         self.last_modified_date = datetime.datetime.now()
         self.creation_date = self.last_modified_date
 
+        # Users who are members of this group.
+        # Note that these links are bidirectional.
+        self.users = set()
+
     def to_json(self):
         return {
             "GroupName": self.group_name,
@@ -221,6 +225,10 @@ class CognitoIdpUser(BaseModel):
         self.create_date = datetime.datetime.utcnow()
         self.last_modified_date = datetime.datetime.utcnow()
 
+        # Groups this user is a member of.
+        # Note that these links are bidirectional.
+        self.groups = set()
+
     def _base_json(self):
         return {
             "UserPoolId": self.user_pool_id,
@@ -428,8 +436,34 @@ class CognitoIdpBackend(BaseBackend):
         if group_name not in user_pool.groups:
             raise ResourceNotFoundError(group_name)
 
+        group = user_pool.groups[group_name]
+        for user in group.users:
+            user.groups.remove(group)
+
         del user_pool.groups[group_name]
 
+    def admin_add_user_to_group(self, user_pool_id, group_name, username):
+        group = self.get_group(user_pool_id, group_name)
+        user = self.admin_get_user(user_pool_id, username)
+
+        group.users.add(user)
+        user.groups.add(group)
+
+    def list_users_in_group(self, user_pool_id, group_name):
+        group = self.get_group(user_pool_id, group_name)
+        return list(group.users)
+
+    def admin_list_groups_for_user(self, user_pool_id, username):
+        user = self.admin_get_user(user_pool_id, username)
+        return list(user.groups)
+
+    def admin_remove_user_from_group(self, user_pool_id, group_name, username):
+        group = self.get_group(user_pool_id, group_name)
+        user = self.admin_get_user(user_pool_id, username)
+
+        group.users.discard(user)
+        user.groups.discard(group)
+
     # User
     def admin_create_user(self, user_pool_id, username, temporary_password, attributes):
         user_pool = self.user_pools.get(user_pool_id)
@@ -465,6 +499,10 @@ class CognitoIdpBackend(BaseBackend):
         if username not in user_pool.users:
             raise ResourceNotFoundError(username)
 
+        user = user_pool.users[username]
+        for group in user.groups:
+            group.users.remove(user)
+
         del user_pool.users[username]
 
     def _log_user_in(self, user_pool, client, username):

moto/cognitoidp/responses.py

@@ -170,6 +170,48 @@ class CognitoIdpResponse(BaseResponse):
         cognitoidp_backends[self.region].delete_group(user_pool_id, group_name)
         return ""
 
+    def admin_add_user_to_group(self):
+        user_pool_id = self._get_param("UserPoolId")
+        username = self._get_param("Username")
+        group_name = self._get_param("GroupName")
+
+        cognitoidp_backends[self.region].admin_add_user_to_group(
+            user_pool_id,
+            group_name,
+            username,
+        )
+
+        return ""
+
+    def list_users_in_group(self):
+        user_pool_id = self._get_param("UserPoolId")
+        group_name = self._get_param("GroupName")
+        users = cognitoidp_backends[self.region].list_users_in_group(user_pool_id, group_name)
+        return json.dumps({
+            "Users": [user.to_json(extended=True) for user in users],
+        })
+
+    def admin_list_groups_for_user(self):
+        username = self._get_param("Username")
+        user_pool_id = self._get_param("UserPoolId")
+        groups = cognitoidp_backends[self.region].admin_list_groups_for_user(user_pool_id, username)
+        return json.dumps({
+            "Groups": [group.to_json() for group in groups],
+        })
+
+    def admin_remove_user_from_group(self):
+        user_pool_id = self._get_param("UserPoolId")
+        username = self._get_param("Username")
+        group_name = self._get_param("GroupName")
+
+        cognitoidp_backends[self.region].admin_remove_user_from_group(
+            user_pool_id,
+            group_name,
+            username,
+        )
+
+        return ""
+
     # User
     def admin_create_user(self):
         user_pool_id = self._get_param("UserPoolId")