Simplify a few unit tests using the USER_INJECT model. (#1173)

This showed that we were flashing 'CODE HAS_BEEN_SENT' even on errors. Fixed that (unified signin).

Author: jwag956

CHANGES.rst

@@ -10,7 +10,7 @@ Released TBD
 
 Features & Improvements
 +++++++++++++++++++++++
-- (:pr:`xx`) Add API :py:meth:`.UserMixin.check_tf_required` to allow applications to control which users
+- (:pr:`1170`) Add API :py:meth:`.UserMixin.check_tf_required` to allow applications to control which users
   require two-factor authentication.
 
 Fixes

flask_security/unified_signin.py

@@ -451,8 +451,9 @@ def us_signin_send_code() -> ResponseValue:
                 form, include_user=False, error_status_code=500 if msg else 200
             )
 
-        # Make sure same response as non-setup method below
-        do_flash(*generic_message("CODE_HAS_BEEN_SENT", "GENERIC_US_SIGNIN"))
+        if not msg:
+            # Make sure same response as non-setup method below
+            do_flash(*generic_message("CODE_HAS_BEEN_SENT", "GENERIC_US_SIGNIN"))
 
         return _security.render_template(
             cv("US_SIGNIN_TEMPLATE"),

tests/test_two_factor.py

@@ -1332,47 +1332,20 @@ def test_bad_sender(app, client, get_message):
     )
 
 
-def test_replace_send_code(app, get_message):
-    pytest.importorskip("sqlalchemy")
-    pytest.importorskip("flask_sqlalchemy")
+cresponse = [None, "That didnt work out as we planned", "Failed Again"]
 
-    # replace tf_send_code - and have it return an error to check that.
-    from flask_sqlalchemy import SQLAlchemy
-    from flask_security.models import fsqla_v2 as fsqla
-    from flask_security import Security, hash_password
 
-    app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:"
-    db = SQLAlchemy(app)
+def _tf_send_code(self, method, totp_secret, phone_number):
+    return self.code_response.pop(0)
 
-    fsqla.FsModels.set_db_info(db)
 
-    class Role(db.Model, fsqla.FsRoleMixin):
-        pass
-
-    class User(db.Model, fsqla.FsUserMixin):
-        rv = [None, "That didnt work out as we planned", "Failed Again"]
-
-        def tf_send_security_token(self, method, **kwargs):
-            return User.rv.pop(0)
-
-    with app.app_context():
-        db.create_all()
-
-    ds = SQLAlchemyUserDatastore(db, User, Role)
-    app.security = Security(app, datastore=ds)
-
-    with app.app_context():
-        client = app.test_client()
-
-        ds.create_user(
-            email="trp@lp.com",
-            password=hash_password("password"),
-            tf_primary_method="sms",
-            tf_totp_secret=app.security._totp_factory.generate_totp_secret(),
-        )
-        ds.commit()
-
-    data = dict(email="trp@lp.com", password="password")
+@pytest.mark.app_settings(
+    TESTING_USER_INJECT=dict(
+        tf_send_security_token=_tf_send_code, code_response=cresponse
+    )
+)
+def test_replace_send_code(app, client, get_message):
+    data = dict(email="gal@lp.com", password="password")
     response = client.post("/login", data=data, follow_redirects=True)
     assert b"Please enter your authentication code" in response.data
     rescue_data = dict(help_setup="email")
@@ -1384,8 +1357,6 @@ def tf_send_security_token(self, method, **kwargs):
     response = client.post("/tf-rescue", json=rescue_data, headers=headers)
     assert response.status_code == 500
     assert response.json["response"]["field_errors"]["help_setup"][0] == "Failed Again"
-    with app.app_context():
-        db.engine.dispose()
 
 
 def test_propagate_next(app, client):

tests/test_unified_signin.py

@@ -41,7 +41,6 @@
 
 from flask_security import (
     SmsSenderFactory,
-    SQLAlchemyUserDatastore,
     UserMixin,
     uia_email_mapper,
     uia_phone_mapper,
@@ -1680,49 +1679,18 @@ def test_bad_sender(app, client, get_message):
     ) == get_message("FAILED_TO_SEND_CODE")
 
 
-@pytest.mark.registerable()
-def test_replace_send_code(app, get_message):
-    pytest.importorskip("sqlalchemy")
-    pytest.importorskip("flask_sqlalchemy")
-
-    from flask_sqlalchemy import SQLAlchemy
-    from flask_security.models import fsqla_v2 as fsqla
-    from flask_security import Security, us_send_security_token
-
-    app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:"
-    db = SQLAlchemy(app)
-
-    fsqla.FsModels.set_db_info(db)
-
-    class Role(db.Model, fsqla.FsRoleMixin):
-        pass
-
-    class User(db.Model, fsqla.FsUserMixin):
-        def us_send_security_token(self, method, **kwargs):
-            assert method == "sms"
-            us_send_security_token(self, method, **kwargs)
-
-    with app.app_context():
-        db.create_all()
-
-    ds = SQLAlchemyUserDatastore(db, User, Role)
-    app.security = Security(app, datastore=ds)
+def _send_code(self, method, **kwargs):
+    assert method == "sms"
+    return "NO SMS AVAILABLE"
 
-    client = app.test_client()
 
-    # since we don't use client fixture - have to add user
-    data = dict(email="trp@lp.com", password="password", password_confirm="password")
-    response = client.post("/register", data=data, follow_redirects=True)
-    assert b"Welcome trp@lp.com" in response.data
-    logout(client)
-
-    set_phone(app, email="trp@lp.com")
-    data = dict(identity="trp@lp.com", chosen_method="sms")
+@pytest.mark.app_settings(TESTING_USER_INJECT=dict(us_send_security_token=_send_code))
+@pytest.mark.registerable()
+def test_replace_send_code(app, client, get_message):
+    set_phone(app, "gal@lp.com")
+    data = dict(identity="gal@lp.com", chosen_method="sms")
     response = client.post("/us-signin/send-code", data=data, follow_redirects=True)
-    assert b"Code has been sent" in response.data
-
-    with app.app_context():
-        db.engine.dispose()  # sqlite wants everything cleaned up
+    assert b"NO SMS AVAILABLE" in response.data
 
 
 @pytest.mark.settings(us_enabled_methods=["password"])