This is an automated email from the ASF dual-hosted git repository. ephraimanierobi pushed a commit to branch v2-3-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 739539a0242c67ad80f2aa90bce0dc3fd782d022 Author: blag <[email protected]> AuthorDate: Thu Apr 28 11:59:24 2022 -0700 Fix update user auth stats (#23314) (cherry picked from commit 1c886eaf99d9044cdf90bf7900dc489a4d40db79) --- airflow/www/fab_security/manager.py | 13 ++++--- tests/www/test_security.py | 67 +++++++++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 5 deletions(-) diff --git a/airflow/www/fab_security/manager.py b/airflow/www/fab_security/manager.py index 5e32b23ac9..e34a3f736a 100644 --- a/airflow/www/fab_security/manager.py +++ b/airflow/www/fab_security/manager.py @@ -822,12 +822,15 @@ class BaseSecurityManager: def update_user_auth_stat(self, user, success=True): """ - Update authentication successful to user. - + Update user authentication stats upon successful/unsuccessful + authentication attempts. :param user: - The authenticated user model + The identified (but possibly not successfully authenticated) user + model :param success: - Default to true, if false increments fail_login_count on user model + Defaults to true, if true increments login_count, updates + last_login, and resets fail_login_count to 0, if false increments + fail_login_count on user model. """ if not user.login_count: user.login_count = 0 @@ -835,10 +838,10 @@ class BaseSecurityManager: user.fail_login_count = 0 if success: user.login_count += 1 + user.last_login = datetime.datetime.now() user.fail_login_count = 0 else: user.fail_login_count += 1 - user.last_login = datetime.datetime.now() self.update_user(user) def auth_user_db(self, username, password): diff --git a/tests/www/test_security.py b/tests/www/test_security.py index e1da8c0d54..8c90062600 100644 --- a/tests/www/test_security.py +++ b/tests/www/test_security.py @@ -16,12 +16,14 @@ # specific language governing permissions and limitations # under the License. import contextlib +import datetime import logging from unittest import mock import pytest from flask_appbuilder import SQLA, Model, expose, has_access from flask_appbuilder.views import BaseView, ModelView +from freezegun import freeze_time from sqlalchemy import Column, Date, Float, Integer, String from airflow.exceptions import AirflowException @@ -814,3 +816,68 @@ def test_fab_models_use_airflow_base_meta(): # TODO: move this test to appropriate place when we have more tests for FAB models user = User() assert user.metadata is Base.metadata + + [email protected]() +def mock_security_manager(app_builder): + mocked_security_manager = MockSecurityManager(appbuilder=app_builder) + mocked_security_manager.update_user = mock.MagicMock() + return mocked_security_manager + + [email protected]() +def new_user(): + user = mock.MagicMock() + user.login_count = None + user.fail_login_count = None + user.last_login = None + return user + + [email protected]() +def old_user(): + user = mock.MagicMock() + user.login_count = 42 + user.fail_login_count = 9 + user.last_login = datetime.datetime(1984, 12, 1, 0, 0, 0) + return user + + +@freeze_time(datetime.datetime(1985, 11, 5, 1, 24, 0)) # Get the Delorean, doc! +def test_update_user_auth_stat_first_successful_auth(mock_security_manager, new_user): + mock_security_manager.update_user_auth_stat(new_user, success=True) + + assert new_user.login_count == 1 + assert new_user.fail_login_count == 0 + assert new_user.last_login == datetime.datetime(1985, 11, 5, 1, 24, 0) + assert mock_security_manager.update_user.called_once + + +@freeze_time(datetime.datetime(1985, 11, 5, 1, 24, 0)) +def test_update_user_auth_stat_subsequent_successful_auth(mock_security_manager, old_user): + mock_security_manager.update_user_auth_stat(old_user, success=True) + + assert old_user.login_count == 43 + assert old_user.fail_login_count == 0 + assert old_user.last_login == datetime.datetime(1985, 11, 5, 1, 24, 0) + assert mock_security_manager.update_user.called_once + + +@freeze_time(datetime.datetime(1985, 11, 5, 1, 24, 0)) +def test_update_user_auth_stat_first_unsuccessful_auth(mock_security_manager, new_user): + mock_security_manager.update_user_auth_stat(new_user, success=False) + + assert new_user.login_count == 0 + assert new_user.fail_login_count == 1 + assert new_user.last_login is None + assert mock_security_manager.update_user.called_once + + +@freeze_time(datetime.datetime(1985, 11, 5, 1, 24, 0)) +def test_update_user_auth_stat_subsequent_unsuccessful_auth(mock_security_manager, old_user): + mock_security_manager.update_user_auth_stat(old_user, success=False) + + assert old_user.login_count == 42 + assert old_user.fail_login_count == 10 + assert old_user.last_login == datetime.datetime(1984, 12, 1, 0, 0, 0) + assert mock_security_manager.update_user.called_once
