fix(user-dao): return user model instances (#30020)

superset/daos/user.py

@@ -21,7 +21,7 @@ import logging
 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 @@ logger = logging.getLogger(__name__)
 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:

tests/unit_tests/dao/user_test.py

@@ -14,79 +14,65 @@
 # 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 test_get_by_id_found(admin_user: User, after_each: None) -> None:  # noqa: F811
-def mock_db_session(mocker):
+    user = UserDAO.get_by_id(admin_user.id)
-    db = mocker.patch("superset.daos.user.db", autospec=True)
+    assert user.id == admin_user.id
-    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):
+def test_get_by_id_not_found():
-    # 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_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
     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):
+def test_set_avatar_url_with_existing_attributes(
-    # Setup
+    admin_user: User,  # noqa: F811
-    user = User()
+    after_each: None,  # noqa: F811
-    user.id = 1
+) -> None:
-    user.extra_attributes = [UserAttribute(user_id=user.id, avatar_url="old_url")]
+    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)
+    UserDAO.set_avatar_url(admin_user, new_url)
-
+    user = UserDAO.get_by_id(admin_user.id)
-    # Assert
     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):
+def test_set_avatar_url_without_existing_attributes(
-    # Setup
+    admin_user: User,  # noqa: F811
-    user = User()
+    after_each: None,  # noqa: F811
-    user.id = 1
+) -> None:
-    user.extra_attributes = []
-
-    # Execute
     new_url = "http://newurl.com"
-    UserDAO.set_avatar_url(user, new_url)
+    UserDAO.set_avatar_url(admin_user, new_url)
 
-    # Assert
+    user = UserDAO.get_by_id(admin_user.id)
-    assert len(user.extra_attributes) == 1
+    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"
+
+    monkeypatch.setattr(security_manager, "user_model", CustomUserModel)
+
+    user = UserDAO.get_by_id(admin_user.id)
+    assert isinstance(user, CustomUserModel)