diff --git a/tests/test_setup_superuser.py b/tests/test_setup_superuser.py index 5c05e8ba7..f1566d9ae 100644 --- a/tests/test_setup_superuser.py +++ b/tests/test_setup_superuser.py @@ -1,4 +1,4 @@ -from unittest.mock import patch, Mock, MagicMock +from unittest.mock import patch, Mock, MagicMock, call from langflow.services.database.models.user.user import User from langflow.services.settings.constants import ( DEFAULT_SUPERUSER, @@ -14,12 +14,30 @@ def test_setup_superuser( mock_get_session, mock_create_super_user, mock_get_settings_manager ): # Test when AUTO_LOGIN is True + calls = [] mock_settings_manager = Mock() mock_settings_manager.auth_settings.AUTO_LOGIN = True + mock_settings_manager.auth_settings.SUPERUSER = DEFAULT_SUPERUSER + mock_settings_manager.auth_settings.SUPERUSER_PASSWORD = DEFAULT_SUPERUSER_PASSWORD mock_get_settings_manager.return_value = mock_settings_manager + mock_session = Mock() + mock_session.query.return_value.filter.return_value.first.return_value = ( + mock_session + ) + # return value of get_session is a generator + mock_get_session.return_value = iter([mock_session, mock_session, mock_session]) setup_superuser() - mock_get_session.assert_not_called() - mock_create_super_user.assert_not_called() + mock_session.query.assert_called_once_with(User) + actual_expr = mock_session.query.return_value.filter.call_args[0][0] + expected_expr = User.username == DEFAULT_SUPERUSER + + assert str(actual_expr) == str(expected_expr) + create_call = call( + db=mock_session, username=DEFAULT_SUPERUSER, password=DEFAULT_SUPERUSER_PASSWORD + ) + calls.append(create_call) + mock_create_super_user.assert_has_calls(calls) + assert 1 == mock_create_super_user.call_count def reset_mock_credentials(): mock_settings_manager.auth_settings.SUPERUSER = DEFAULT_SUPERUSER @@ -38,21 +56,17 @@ def test_setup_superuser( ) mock_get_settings_manager.return_value = mock_settings_manager - mock_session = Mock() - mock_session.query.return_value.filter.return_value.first.return_value = ( - mock_session - ) - # return value of get_session is a generator - mock_get_session.return_value = iter([mock_session, mock_session]) + setup_superuser() - mock_session.query.assert_called_once_with(User) + mock_session.query.assert_called_with(User) actual_expr = mock_session.query.return_value.filter.call_args[0][0] expected_expr = User.username == ADMIN_USER_NAME assert str(actual_expr) == str(expected_expr) - mock_create_super_user.assert_called_once_with( - db=mock_session, username=ADMIN_USER_NAME, password="password" - ) + create_call = call(db=mock_session, username=ADMIN_USER_NAME, password="password") + calls.append(create_call) + mock_create_super_user.assert_has_calls(calls) + assert 2 == mock_create_super_user.call_count # Test that superuser credentials are reset mock_settings_manager.auth_settings.reset_credentials.assert_called_once() assert mock_settings_manager.auth_settings.SUPERUSER != ADMIN_USER_NAME @@ -72,9 +86,6 @@ def test_setup_superuser( assert str(actual_expr) == str(expected_expr) - # Called once in the previous test - mock_create_super_user.assert_called_once() - @patch("langflow.services.utils.get_settings_manager") @patch("langflow.services.utils.get_session")