From aeecd6279fbc98a38690721b1c69deeaaabaedfc Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 18 Sep 2023 17:43:31 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=80=20feat(test=5Fsetup=5Fsuperuser.py?= =?UTF-8?q?):=20add=20tests=20for=20setup=5Fsuperuser=20and=20teardown=5Fs?= =?UTF-8?q?uperuser=20functions?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 🔒 test(test_setup_superuser.py): test when AUTO_LOGIN is True, setup_superuser should not call get_session and create_super_user 🔒 test(test_setup_superuser.py): test when username and password are default, setup_superuser should call get_session, query, filter, and create_super_user with correct arguments, and reset superuser credentials 🔒 test(test_setup_superuser.py): test when superuser already exists, setup_superuser should call get_session, query, and filter with correct arguments, and not call create_super_user 🔒 test(test_teardown_superuser_default_superuser): test teardown_superuser with default superuser, should call get_session, query, filter, delete, and commit with correct arguments 🔒 test(test_teardown_superuser_no_default_superuser): test teardown_superuser without default superuser, should not call get_session, query, delete, and commit --- tests/test_setup_superuser.py | 127 ++++++++++++++++++++++++++++++++++ 1 file changed, 127 insertions(+) create mode 100644 tests/test_setup_superuser.py diff --git a/tests/test_setup_superuser.py b/tests/test_setup_superuser.py new file mode 100644 index 000000000..202910737 --- /dev/null +++ b/tests/test_setup_superuser.py @@ -0,0 +1,127 @@ +from unittest.mock import patch, Mock, MagicMock +from langflow.services.database.models.user.user import User +from langflow.services.settings.constants import ( + DEFAULT_SUPERUSER, + DEFAULT_SUPERUSER_PASSWORD, +) +from langflow.services.utils import setup_superuser, teardown_superuser + + +@patch("langflow.services.utils.get_settings_manager") +@patch("langflow.services.utils.create_super_user") +@patch("langflow.services.utils.get_session") +def test_setup_superuser( + mock_get_session, mock_create_super_user, mock_get_settings_manager +): + # Test when AUTO_LOGIN is True + mock_settings_manager = Mock() + mock_settings_manager.auth_settings.AUTO_LOGIN = True + mock_get_settings_manager.return_value = mock_settings_manager + setup_superuser() + mock_get_session.assert_not_called() + mock_create_super_user.assert_not_called() + + def reset_mock_credentials(): + mock_settings_manager.auth_settings.SUPERUSER = DEFAULT_SUPERUSER + mock_settings_manager.auth_settings.SUPERUSER_PASSWORD = ( + DEFAULT_SUPERUSER_PASSWORD + ) + + # Test when username and password are default + mock_settings_manager.auth_settings = Mock() + mock_settings_manager.auth_settings.AUTO_LOGIN = False + mock_settings_manager.auth_settings.SUPERUSER = "admin" + mock_settings_manager.auth_settings.SUPERUSER_PASSWORD = "password" + mock_settings_manager.auth_settings.reset_credentials = Mock( + side_effect=reset_mock_credentials + ) + + 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) + actual_expr = mock_session.query.return_value.filter.call_args[0][0] + expected_expr = User.username == "admin" + + assert str(actual_expr) == str(expected_expr) + mock_create_super_user.assert_called_once_with( + db=mock_session, username="admin", password="password" + ) + # Test that superuser credentials are reset + mock_settings_manager.auth_settings.reset_credentials.assert_called_once() + assert mock_settings_manager.auth_settings.SUPERUSER != "admin" + assert mock_settings_manager.auth_settings.SUPERUSER_PASSWORD != "password" + + # Test when superuser already exists + mock_settings_manager.auth_settings.AUTO_LOGIN = False + mock_settings_manager.auth_settings.SUPERUSER = "admin" + mock_settings_manager.auth_settings.SUPERUSER_PASSWORD = "password" + mock_user = Mock() + mock_user.is_superuser = True + mock_session.query.return_value.filter.return_value.first.return_value = mock_user + setup_superuser() + mock_session.query.assert_called_with(User) + actual_expr = mock_session.query.return_value.filter.call_args[0][0] + expected_expr = User.username == "admin" + + 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") +def test_teardown_superuser_default_superuser( + mock_get_session, mock_get_settings_manager +): + mock_settings_manager = MagicMock() + 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 = MagicMock() + mock_user = MagicMock() + mock_user.is_superuser = True + mock_session.query.return_value.filter.return_value.first.return_value = mock_user + mock_get_session.return_value = iter([mock_session]) + + teardown_superuser() + + 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) + mock_session.delete.assert_called_once_with(mock_user) + mock_session.commit.assert_called_once() + + +@patch("langflow.services.utils.get_settings_manager") +@patch("langflow.services.utils.get_session") +def test_teardown_superuser_no_default_superuser( + mock_get_session, mock_get_settings_manager +): + mock_settings_manager = MagicMock() + mock_settings_manager.auth_settings.AUTO_LOGIN = False + mock_settings_manager.auth_settings.SUPERUSER = "admin" + mock_settings_manager.auth_settings.SUPERUSER_PASSWORD = "password" + mock_get_settings_manager.return_value = mock_settings_manager + + mock_session = MagicMock() + mock_user = MagicMock() + mock_user.is_superuser = False + mock_session.query.return_value.filter.return_value.first.return_value = mock_user + mock_get_session.return_value = [mock_session] + + teardown_superuser() + + mock_session.query.assert_not_called() + mock_session.delete.assert_not_called() + mock_session.commit.assert_not_called()