From c3f9433612d41a25fa682b3949fcf75a62c73b2a Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Tue, 5 Dec 2023 18:35:14 -0300 Subject: [PATCH] Fix superuser deletion bug --- src/backend/langflow/services/utils.py | 2 +- tests/conftest.py | 10 ++++------ tests/test_setup_superuser.py | 7 ++++--- 3 files changed, 9 insertions(+), 10 deletions(-) diff --git a/src/backend/langflow/services/utils.py b/src/backend/langflow/services/utils.py index 90fa9893d..8f44b114d 100644 --- a/src/backend/langflow/services/utils.py +++ b/src/backend/langflow/services/utils.py @@ -134,7 +134,7 @@ def teardown_superuser(settings_service, session): from langflow.services.database.models.user.model import User user = session.exec(select(User).where(User.username == username)).first() - if user and user.is_superuser: + if user and user.is_superuser is True: session.delete(user) session.commit() logger.debug("Default superuser removed successfully.") diff --git a/tests/conftest.py b/tests/conftest.py index 7461b1f0f..770a01918 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,4 @@ import json - # we need to import tmpdir import tempfile from contextlib import contextmanager, suppress @@ -10,16 +9,15 @@ import orjson import pytest from fastapi.testclient import TestClient from httpx import AsyncClient -from sqlmodel import Session, SQLModel, create_engine -from sqlmodel.pool import StaticPool -from typer.testing import CliRunner - from langflow.graph.graph.base import Graph from langflow.services.auth.utils import get_password_hash from langflow.services.database.models.flow.model import Flow, FlowCreate from langflow.services.database.models.user.model import User, UserCreate from langflow.services.database.utils import session_getter from langflow.services.deps import get_db_service +from sqlmodel import Session, SQLModel, create_engine +from sqlmodel.pool import StaticPool +from typer.testing import CliRunner if TYPE_CHECKING: from langflow.services.database.service import DatabaseService @@ -305,7 +303,7 @@ def added_vector_store(client, json_vector_store, logged_in_headers): @pytest.fixture def test_component_code(): - path = Path(__file__).parent.absolute() / "data" / "component_test.py" + path = Path(__file__).parent.absolute() / "data" / "component.py" # load the content as a string with open(path, "r") as f: return f.read() diff --git a/tests/test_setup_superuser.py b/tests/test_setup_superuser.py index d0c9abbbd..e8951258f 100644 --- a/tests/test_setup_superuser.py +++ b/tests/test_setup_superuser.py @@ -1,6 +1,7 @@ from unittest.mock import MagicMock, patch -from langflow.services.settings.constants import DEFAULT_SUPERUSER, DEFAULT_SUPERUSER_PASSWORD +from langflow.services.settings.constants import (DEFAULT_SUPERUSER, + DEFAULT_SUPERUSER_PASSWORD) from langflow.services.utils import teardown_superuser # @patch("langflow.services.deps.get_session") @@ -119,11 +120,11 @@ def test_teardown_superuser_no_default_superuser(mock_get_session, mock_get_sett 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_session.exec.return_value.filter.return_value.first.return_value = mock_user mock_get_session.return_value = [mock_session] teardown_superuser(mock_settings_service, mock_session) - mock_session.query.assert_called_once() + mock_session.exec.assert_called_once() mock_session.delete.assert_not_called() mock_session.commit.assert_not_called()