Add implementation for cognitoidp admin_set_user_mfa_preference (#4884)

moto/cognitoidp/models.py

@@ -652,6 +652,7 @@ class CognitoIdpUser(BaseModel):
         self.software_token_mfa_enabled = False
         self.token_verified = False
         self.confirmation_code = None
+        self.preferred_mfa_setting = None
 
         # Groups this user is a member of.
         # Note that these links are bidirectional.
@@ -690,6 +691,7 @@ class CognitoIdpUser(BaseModel):
                     attributes_key: attrs,
                     "MFAOptions": [],
                     "UserMFASettingList": user_mfa_setting_list,
+                    "PreferredMfaSetting": self.preferred_mfa_setting or "",
                 }
             )
 
@@ -1635,7 +1637,9 @@ class CognitoIdpBackend(BaseBackend):
                 _, username = user_pool.access_tokens[access_token]
                 user = self.admin_get_user(user_pool.id, username)
 
-                if software_token_mfa_settings["Enabled"]:
+                if software_token_mfa_settings and software_token_mfa_settings.get(
+                    "Enabled"
+                ):
                     if user.token_verified:
                         user.software_token_mfa_enabled = True
                     else:
@@ -1643,13 +1647,39 @@ class CognitoIdpBackend(BaseBackend):
                             "User has not verified software token mfa"
                         )
 
-                elif sms_mfa_settings["Enabled"]:
+                    if software_token_mfa_settings.get("PreferredMfa"):
+                        user.preferred_mfa_setting = "SOFTWARE_TOKEN_MFA"
+                elif sms_mfa_settings and sms_mfa_settings["Enabled"]:
                     user.sms_mfa_enabled = True
 
+                    if sms_mfa_settings.get("PreferredMfa"):
+                        user.preferred_mfa_setting = "SMS_MFA"
                 return None
         else:
             raise NotAuthorizedError(access_token)
 
+    def admin_set_user_mfa_preference(
+        self, user_pool_id, username, software_token_mfa_settings, sms_mfa_settings
+    ):
+        user = self.admin_get_user(user_pool_id, username)
+
+        if software_token_mfa_settings and software_token_mfa_settings.get("Enabled"):
+            if user.token_verified:
+                user.software_token_mfa_enabled = True
+            else:
+                raise InvalidParameterException(
+                    "User has not verified software token mfa"
+                )
+
+            if software_token_mfa_settings.get("PreferredMfa"):
+                user.preferred_mfa_setting = "SOFTWARE_TOKEN_MFA"
+        elif sms_mfa_settings and sms_mfa_settings.get("Enabled"):
+            user.sms_mfa_enabled = True
+
+            if sms_mfa_settings.get("PreferredMfa"):
+                user.preferred_mfa_setting = "SMS_MFA"
+        return None
+
     def admin_set_user_password(self, user_pool_id, username, password, permanent):
         user = self.admin_get_user(user_pool_id, username)
         user.password = password

moto/cognitoidp/responses.py

@@ -571,6 +571,16 @@ class CognitoIdpResponse(BaseResponse):
         )
         return ""
 
+    def admin_set_user_mfa_preference(self):
+        user_pool_id = self._get_param("UserPoolId")
+        username = self._get_param("Username")
+        software_token_mfa_settings = self._get_param("SoftwareTokenMfaSettings")
+        sms_mfa_settings = self._get_param("SMSMfaSettings")
+        cognitoidp_backends[self.region].admin_set_user_mfa_preference(
+            user_pool_id, username, software_token_mfa_settings, sms_mfa_settings
+        )
+        return ""
+
     def admin_set_user_password(self):
         user_pool_id = self._get_param("UserPoolId")
         username = self._get_param("Username")

tests/test_cognitoidp/test_cognitoidp.py

@@ -3628,6 +3628,7 @@ def test_setting_mfa():
         )
 
         result["UserMFASettingList"].should.have.length_of(1)
+        result["PreferredMfaSetting"].should.equal("SOFTWARE_TOKEN_MFA")
 
 
 @mock_cognitoidp
@@ -3650,6 +3651,44 @@ def test_setting_mfa_when_token_not_verified():
         caught.should.be.true
 
 
+@mock_cognitoidp
+def test_admin_setting_mfa():
+    conn = boto3.client("cognito-idp", "us-west-2")
+
+    user_pool_id = conn.create_user_pool(
+        PoolName=str(uuid.uuid4()), UsernameAttributes=["email"]
+    )["UserPool"]["Id"]
+    username = "test@example.com"
+    conn.admin_create_user(UserPoolId=user_pool_id, Username=username)
+
+    conn.admin_set_user_mfa_preference(
+        Username=username,
+        UserPoolId=user_pool_id,
+        SMSMfaSettings={"Enabled": True, "PreferredMfa": True},
+    )
+    result = conn.admin_get_user(UserPoolId=user_pool_id, Username=username)
+    result["UserMFASettingList"].should.have.length_of(1)
+    result["PreferredMfaSetting"].should.equal("SMS_MFA")
+
+
+@mock_cognitoidp
+def test_admin_setting_mfa_when_token_not_verified():
+    conn = boto3.client("cognito-idp", "us-west-2")
+
+    user_pool_id = conn.create_user_pool(
+        PoolName=str(uuid.uuid4()), UsernameAttributes=["email"]
+    )["UserPool"]["Id"]
+    username = "test@example.com"
+    conn.admin_create_user(UserPoolId=user_pool_id, Username=username)
+
+    with pytest.raises(conn.exceptions.InvalidParameterException):
+        conn.admin_set_user_mfa_preference(
+            Username=username,
+            UserPoolId=user_pool_id,
+            SoftwareTokenMfaSettings={"Enabled": True, "PreferredMfa": True},
+        )
+
+
 @mock_cognitoidp
 def test_respond_to_auth_challenge_with_invalid_secret_hash():
     conn = boto3.client("cognito-idp", "us-west-2")