test: use SQLite sessions in core tools (#39094)

api/tests/unit_tests/core/tools/workflow_as_tool/test_provider.py

@@ -1,12 +1,17 @@
 from __future__ import annotations
 
 import json
+import uuid
+from collections.abc import Iterator
 from types import SimpleNamespace
 from typing import Any, cast
-from unittest.mock import MagicMock, Mock, patch
+from unittest.mock import PropertyMock, patch
 
 import pytest
+from sqlalchemy import Engine
+from sqlalchemy.orm import Session, sessionmaker
 
+from core.db.session_factory import session_factory
 from core.tools.__base.tool_runtime import ToolRuntime
 from core.tools.entities.common_entities import I18nObject
 from core.tools.entities.tool_entities import (
@@ -20,14 +25,31 @@ from core.tools.entities.tool_entities import (
 )
 from core.tools.workflow_as_tool.provider import WorkflowToolProviderController
 from core.tools.workflow_as_tool.tool import WorkflowTool
+from extensions.ext_database import db
 from graphon.variables.input_entities import VariableEntity, VariableEntityType
 from models.account import Account
-from models.model import App
+from models.base import TypeBase
+from models.model import App, AppMode, IconType
 from models.tools import WorkflowToolProvider
 from models.workflow import Workflow, WorkflowType
 
 
-def _controller() -> WorkflowToolProviderController:
+@pytest.fixture
+def database_session(sqlite_engine: Engine) -> Iterator[Session]:
+    models = (Account, App, Workflow, WorkflowToolProvider)
+    tables = [model.metadata.tables[model.__tablename__] for model in models]
+    TypeBase.metadata.create_all(sqlite_engine, tables=tables)
+    session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False)
+
+    with (
+        patch.object(session_factory, "create_session", session_maker),
+        patch.object(type(db), "engine", new_callable=PropertyMock, return_value=sqlite_engine),
+    ):
+        with session_maker() as session:
+            yield session
+
+
+def _controller(provider_id: str = "provider-1") -> WorkflowToolProviderController:
     entity = ToolProviderEntity(
         identity=ToolProviderIdentity(
             author="author",
@@ -38,47 +60,64 @@ def _controller() -> WorkflowToolProviderController:
         ),
         credentials_schema=[],
     )
-    return WorkflowToolProviderController(entity=entity, provider_id="provider-1")
+    return WorkflowToolProviderController(entity=entity, provider_id=provider_id)
 
 
-def _app() -> App:
-    return App(id="app-1")
+def _app(*, tenant_id: str | None = None) -> App:
+    return App(
+        id=str(uuid.uuid4()),
+        tenant_id=tenant_id or str(uuid.uuid4()),
+        name="Workflow App",
+        mode=AppMode.WORKFLOW,
+        icon_type=IconType.EMOJI,
+        icon="workflow",
+        icon_background="#FFFFFF",
+        enable_site=True,
+        enable_api=False,
+    )
 
 
 def _account() -> Account:
     return Account(name="Alice", email="alice@example.com")
 
 
-def _workflow() -> Workflow:
+def _workflow(app: App, account: Account | None = None) -> Workflow:
     return Workflow.new(
-        tenant_id="tenant-1",
-        app_id="app-1",
+        tenant_id=app.tenant_id,
+        app_id=app.id,
         type=WorkflowType.WORKFLOW.value,
         version="1",
         graph=json.dumps({"nodes": []}),
         features="{}",
-        created_by="user-1",
+        created_by=account.id if account else str(uuid.uuid4()),
         environment_variables=[],
         conversation_variables=[],
         rag_pipeline_variables=[],
     )
 
 
-def _db_provider(*, parameter_configuration: str = "[]") -> WorkflowToolProvider:
+def _db_provider(
+    app: App,
+    account: Account,
+    *,
+    parameter_configuration: str = "[]",
+) -> WorkflowToolProvider:
     return WorkflowToolProvider(
         name="workflow_tool",
         label="WF Provider",
         icon="icon.svg",
-        app_id="app-1",
+        app_id=app.id,
         version="1",
-        user_id="user-1",
-        tenant_id="tenant-1",
+        user_id=account.id,
+        tenant_id=app.tenant_id,
         description="desc",
         parameter_configuration=parameter_configuration,
     )
 
 
-def _workflow_tool(name: str = "workflow_tool") -> WorkflowTool:
+def _workflow_tool(name: str = "workflow_tool", *, tenant_id: str | None = None) -> WorkflowTool:
+    app = _app(tenant_id=tenant_id)
+    workflow = _workflow(app)
     return WorkflowTool(
         workflow_as_tool_id="provider-1",
         entity=ToolEntity(
@@ -91,38 +130,46 @@ def _workflow_tool(name: str = "workflow_tool") -> WorkflowTool:
             description=ToolDescription(human=I18nObject(en_US="desc"), llm="desc"),
             parameters=[],
         ),
-        runtime=ToolRuntime(tenant_id="tenant-1"),
-        workflow_app_id="app-1",
-        workflow_entities={"app": _app(), "workflow": _workflow()},
+        runtime=ToolRuntime(tenant_id=app.tenant_id),
+        workflow_app_id=app.id,
+        workflow_entities={"app": app, "workflow": workflow},
         version="1",
         workflow_call_depth=0,
     )
 
 
-def _mock_session_with_begin() -> Mock:
-    session = Mock()
-    begin_cm = Mock()
-    begin_cm.__enter__ = Mock(return_value=None)
-    begin_cm.__exit__ = Mock(return_value=False)
-    session.begin.return_value = begin_cm
-    return session
-
-
-def test_get_db_provider_tool_builds_entity():
-    controller = _controller()
-    session = Mock()
-    workflow = _workflow()
-    session.scalar.return_value = workflow
+def _persist_provider_graph(
+    session: Session,
+    *,
+    parameter_configuration: str = "[]",
+    include_app: bool = True,
+    include_workflow: bool = True,
+) -> tuple[WorkflowToolProvider, App, Account, Workflow]:
+    account = _account()
     app = _app()
-    db_provider = _db_provider(
+    workflow = _workflow(app, account)
+    db_provider = _db_provider(app, account, parameter_configuration=parameter_configuration)
+
+    session.add_all([account, db_provider])
+    if include_app:
+        session.add(app)
+    if include_workflow:
+        session.add(workflow)
+    session.commit()
+    return db_provider, app, account, workflow
+
+
+def test_get_db_provider_tool_builds_entity(database_session: Session):
+    db_provider, app, user, _ = _persist_provider_graph(
+        database_session,
         parameter_configuration=json.dumps(
             [
                 {"name": "country", "description": "Country", "form": ToolParameter.ToolParameterForm.FORM.value},
                 {"name": "files", "description": "files", "form": ToolParameter.ToolParameterForm.FORM.value},
             ]
-        )
+        ),
     )
-    user = _account()
+    controller = _controller(db_provider.id)
     variables = [
         VariableEntity(
             variable="country",
@@ -152,7 +199,7 @@ def test_get_db_provider_tool_builds_entity():
             return_value=outputs,
         ),
     ):
-        tool = controller._get_db_provider_tool(db_provider, app, session=session, user=user)
+        tool = controller._get_db_provider_tool(db_provider, app, session=database_session, user=user)
 
     assert tool.entity.identity.name == "workflow_tool"
     # "json" output is reserved for ToolInvokeMessage.VariableMessage and filtered out.
@@ -175,61 +222,55 @@ def test_get_tool_returns_hit_or_none():
 
 def test_get_tools_returns_cached():
     controller = _controller()
-    cached_tools = [_workflow_tool("wf-cached")]
+    cached_tools = [_workflow_tool("wf-cached", tenant_id="tenant-1")]
     controller.tools = cached_tools
 
     assert controller.get_tools("tenant-1") == cached_tools
 
 
-def test_from_db_builds_controller():
-    app = _app()
-    user = _account()
-    db_provider = _db_provider()
-    session = _mock_session_with_begin()
-    session.scalar.return_value = db_provider
-    session.get.side_effect = [app, user]
-    fake_cm = MagicMock()
-    fake_cm.__enter__.return_value = session
-    fake_cm.__exit__.return_value = False
-    fake_session_factory = Mock()
-    fake_session_factory.create_session.return_value = fake_cm
+def test_from_db_builds_controller(database_session: Session):
+    db_provider, app, user, workflow = _persist_provider_graph(database_session)
+
+    with (
+        patch(
+            "core.tools.workflow_as_tool.provider.WorkflowAppConfigManager.convert_features",
+            return_value=SimpleNamespace(file_upload=False),
+        ),
+        patch(
+            "core.tools.workflow_as_tool.provider.WorkflowToolConfigurationUtils.get_workflow_graph_variables",
+            return_value=[],
+        ),
+        patch(
+            "core.tools.workflow_as_tool.provider.WorkflowToolConfigurationUtils.get_workflow_graph_output",
+            return_value=[],
+        ),
+    ):
+        built = WorkflowToolProviderController.from_db(db_provider)
 
-    with patch("core.tools.workflow_as_tool.provider.session_factory", fake_session_factory):
-        with patch.object(
-            WorkflowToolProviderController,
-            "_get_db_provider_tool",
-            return_value=_workflow_tool("wf"),
-        ):
-            built = WorkflowToolProviderController.from_db(db_provider)
     assert isinstance(built, WorkflowToolProviderController)
-    assert built.tools
+    assert built.entity.identity.author == user.name
+    assert built.provider_id == db_provider.id
+    assert built.tools is not None
+    assert built.tools[0].workflow_app_id == app.id
+    assert built.tools[0].workflow_entities["workflow"].id == workflow.id
 
 
-def test_get_tools_returns_empty_when_provider_missing():
-    controller = _controller()
+def test_get_tools_returns_empty_when_provider_missing(database_session: Session):
+    db_provider, _, _, _ = _persist_provider_graph(database_session)
+    controller = _controller(db_provider.id)
     controller.tools = None
 
-    with patch("core.tools.workflow_as_tool.provider.db") as mock_db:
-        mock_db.engine = object()
-        with patch("core.tools.workflow_as_tool.provider.Session") as session_cls:
-            session = _mock_session_with_begin()
-            session.scalar.return_value = None
-            session_cls.return_value.__enter__.return_value = session
-
-            assert controller.get_tools("tenant-1") == []
+    assert controller.get_tools(str(uuid.uuid4())) == []
 
 
-def test_get_tools_raises_when_app_missing():
-    controller = _controller()
+def test_get_tools_raises_when_app_missing(database_session: Session):
+    db_provider, _, _, _ = _persist_provider_graph(
+        database_session,
+        include_app=False,
+        include_workflow=False,
+    )
+    controller = _controller(db_provider.id)
     controller.tools = None
-    db_provider = _db_provider()
 
-    with patch("core.tools.workflow_as_tool.provider.db") as mock_db:
-        mock_db.engine = object()
-        with patch("core.tools.workflow_as_tool.provider.Session") as session_cls:
-            session = _mock_session_with_begin()
-            session.scalar.return_value = db_provider
-            session.get.return_value = None
-            session_cls.return_value.__enter__.return_value = session
-            with pytest.raises(ValueError, match="app not found"):
-                controller.get_tools("tenant-1")
+    with pytest.raises(ValueError, match="app not found"):
+        controller.get_tools(db_provider.tenant_id)