From fcf04502949b58fbcd7225ec7d10e9c73ae316d5 Mon Sep 17 00:00:00 2001
From: Ville Brofeldt <33317356+villebro@users.noreply.github.com>
Date: Mon, 26 Aug 2024 16:50:15 -0700
Subject: [PATCH] fix(user-dao): return user model instances (#30020)

---
 superset/daos/user.py             |  4 +-
 tests/unit_tests/dao/user_test.py | 94 +++++++++++++------------------
 2 files changed, 42 insertions(+), 56 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..6066b0e7dfc5d 100644
--- a/tests/unit_tests/dao/user_test.py
+++ b/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 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(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_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_not_found():
+    with pytest.raises(NoResultFound):
+        UserDAO.get_by_id(123456)
 
-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)
+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()
 
+    new_url = "http://newurl.com"
+    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
 
-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")]
 
-    # 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
+    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_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 = []
+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"
 
-    # Execute
-    new_url = "http://newurl.com"
-    UserDAO.set_avatar_url(user, new_url)
+    monkeypatch.setattr(security_manager, "user_model", CustomUserModel)
 
-    # Assert
-    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
+    user = UserDAO.get_by_id(admin_user.id)
+    assert isinstance(user, CustomUserModel)