From 8332eff9f435dfdaa13eb3e70129cba27b7def2a Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Mon, 26 Aug 2024 13:49:13 -0700 Subject: [PATCH 1/2] fix(user-dao): return user_model instances --- superset/daos/user.py | 4 ++-- tests/unit_tests/dao/user_test.py | 16 ++++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/superset/daos/user.py b/superset/daos/user.py index 90a9b2bd2f6e5..475e3252a6b21 100644 --- a/superset/daos/user.py +++ b/superset/daos/user.py @@ -21,7 +21,7 @@ from flask_appbuilder.security.sqla.models import User from superset.daos.base import BaseDAO -from superset.extensions import db +from superset.extensions import db, security_manager from superset.models.user_attributes import UserAttribute logger = logging.getLogger(__name__) @@ -30,7 +30,7 @@ class UserDAO(BaseDAO[User]): @staticmethod def get_by_id(user_id: int) -> User: - return db.session.query(User).filter_by(id=user_id).one() + return db.session.query(security_manager.user_model).filter_by(id=user_id).one() @staticmethod def set_avatar_url(user: User, url: str) -> None: diff --git a/tests/unit_tests/dao/user_test.py b/tests/unit_tests/dao/user_test.py index bf65c51121fac..7dc8a67b8afb8 100644 --- a/tests/unit_tests/dao/user_test.py +++ b/tests/unit_tests/dao/user_test.py @@ -21,7 +21,9 @@ from sqlalchemy.orm.exc import NoResultFound from superset.daos.user import UserDAO +from superset.extensions import security_manager from superset.models.user_attributes import UserAttribute +from tests.unit_tests.fixtures.common import admin_user, after_each # noqa: F401 @pytest.fixture @@ -90,3 +92,17 @@ def test_set_avatar_url_without_existing_attributes(mock_db_session): assert len(user.extra_attributes) == 1 assert user.extra_attributes[0].avatar_url == new_url mock_db_session.add.assert_called() # New attribute should be added + + +def test_get_by_id_custom_user_class( + monkeypatch: pytest.MonkeyPatch, + admin_user: User, # noqa: F811 + after_each: None, # noqa: F811 +): + class CustomUserModel(User): + __tablename__ = "ab_user" + + monkeypatch.setattr(security_manager, "user_model", CustomUserModel) + + user = UserDAO.get_by_id(admin_user.id) + assert isinstance(user, CustomUserModel) From 6e405f22900d5cb7d453bc975c7964834d382cba Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Mon, 26 Aug 2024 13:59:25 -0700 Subject: [PATCH 2/2] remove mocks --- tests/unit_tests/dao/user_test.py | 82 ++++++++++--------------------- 1 file changed, 26 insertions(+), 56 deletions(-) diff --git a/tests/unit_tests/dao/user_test.py b/tests/unit_tests/dao/user_test.py index 7dc8a67b8afb8..6066b0e7dfc5d 100644 --- a/tests/unit_tests/dao/user_test.py +++ b/tests/unit_tests/dao/user_test.py @@ -14,91 +14,61 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from unittest.mock import MagicMock +from __future__ import annotations import pytest from flask_appbuilder.security.sqla.models import User -from sqlalchemy.orm.exc import NoResultFound +from sqlalchemy.exc import NoResultFound +from superset import db from superset.daos.user import UserDAO from superset.extensions import security_manager from superset.models.user_attributes import UserAttribute from tests.unit_tests.fixtures.common import admin_user, after_each # noqa: F401 -@pytest.fixture -def mock_db_session(mocker): - db = mocker.patch("superset.daos.user.db", autospec=True) - db.session = MagicMock() - db.session.query = MagicMock() - db.session.commit = MagicMock() - db.session.query.return_value = MagicMock() - return db.session - - -def test_get_by_id_found(mock_db_session): - # Setup - user_id = 1 - mock_user = User() - mock_user.id = user_id - mock_query = mock_db_session.query.return_value - mock_query.filter_by.return_value.one.return_value = mock_user - - # Execute - UserDAO.get_by_id(user_id) # noqa: F841 - - # Assert - mock_db_session.query.assert_called_with(User) - mock_query.filter_by.assert_called_with(id=user_id) - +def test_get_by_id_found(admin_user: User, after_each: None) -> None: # noqa: F811 + user = UserDAO.get_by_id(admin_user.id) + assert user.id == admin_user.id -def test_get_by_id_not_found(mock_db_session): - # Setup - user_id = 1 - mock_query = mock_db_session.query.return_value - mock_query.filter_by.return_value.one.side_effect = NoResultFound - # Execute & Assert +def test_get_by_id_not_found(): with pytest.raises(NoResultFound): - UserDAO.get_by_id(user_id) + UserDAO.get_by_id(123456) -def test_set_avatar_url_with_existing_attributes(mock_db_session): - # Setup - user = User() - user.id = 1 - user.extra_attributes = [UserAttribute(user_id=user.id, avatar_url="old_url")] +def test_set_avatar_url_with_existing_attributes( + admin_user: User, # noqa: F811 + after_each: None, # noqa: F811 +) -> None: + admin_user.extra_attributes = [ + UserAttribute(user_id=admin_user.id, avatar_url="old_url"), + ] + db.session.flush() - # Execute new_url = "http://newurl.com" - UserDAO.set_avatar_url(user, new_url) - - # Assert + UserDAO.set_avatar_url(admin_user, new_url) + user = UserDAO.get_by_id(admin_user.id) assert user.extra_attributes[0].avatar_url == new_url - mock_db_session.add.assert_not_called() # No new attributes should be added - -def test_set_avatar_url_without_existing_attributes(mock_db_session): - # Setup - user = User() - user.id = 1 - user.extra_attributes = [] - # Execute +def test_set_avatar_url_without_existing_attributes( + admin_user: User, # noqa: F811 + after_each: None, # noqa: F811 +) -> None: new_url = "http://newurl.com" - UserDAO.set_avatar_url(user, new_url) + UserDAO.set_avatar_url(admin_user, new_url) - # Assert - assert len(user.extra_attributes) == 1 + user = UserDAO.get_by_id(admin_user.id) + assert len(admin_user.extra_attributes) == 1 assert user.extra_attributes[0].avatar_url == new_url - mock_db_session.add.assert_called() # New attribute should be added def test_get_by_id_custom_user_class( monkeypatch: pytest.MonkeyPatch, admin_user: User, # noqa: F811 after_each: None, # noqa: F811 -): +) -> None: class CustomUserModel(User): __tablename__ = "ab_user"