test: migrate application pipelines sessions and ORM models to SQLite (#40510)

This commit is contained in:
Asuka Minato authored and GitHub committed 2026-09-19 09:26:25 +00:00
1 parent ce5dc6f314
commit ceacf275b1
8 files changed
+527 -236

No files matched your search

@@ -1,22 +1,66 @@
import uuid
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
from pytest_mock import MockerFixture
from sqlalchemy.orm import Session
from core.app.app_config.entities import EasyUIBasedAppModelConfigFrom
from core.app.apps.agent_chat.app_config_manager import (
AgentChatAppConfigManager,
)
from core.entities.agent_entities import PlanningStrategy
from models.enums import ConversationFromSource
from models.model import App, AppMode, AppModelConfig, Conversation
def _app() -> App:
return App(
id="app1",
tenant_id="tenant",
name="Agent chat app",
description="",
mode=AppMode.AGENT_CHAT,
enable_site=False,
enable_api=False,
)
def _config() -> AppModelConfig:
return AppModelConfig(app_id="app1")
def _conversation() -> Conversation:
return Conversation(
id="conv",
app_id="app1",
mode=AppMode.AGENT_CHAT,
name="Conversation",
inputs={},
from_source=ConversationFromSource.CONSOLE,
)
_CURRENT_SESSION: Session | None = None
@pytest.fixture(autouse=True)
def _bind_unbound_session(unbound_session: Session):
global _CURRENT_SESSION
_CURRENT_SESSION = unbound_session
yield
_CURRENT_SESSION = None
def _session() -> Session:
assert _CURRENT_SESSION is not None
return _CURRENT_SESSION
class TestAgentChatAppConfigManagerGetAppConfig:
def test_get_app_config_override_config(self, mocker: MockerFixture):
app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent-chat")
app_model_config = mocker.MagicMock(id="cfg1")
app_model_config.to_dict.return_value = {"ignored": True}
app_model = _app()
app_model_config = _config()
override_config = {"model": {"provider": "p"}}
@@ -49,11 +93,15 @@ class TestAgentChatAppConfigManagerGetAppConfig:
assert result.external_data_variables == "external"
def test_get_app_config_conversation_specific(self, mocker: MockerFixture):
app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent-chat")
app_model_config = mocker.MagicMock(id="cfg1")
app_model_config.to_dict.return_value = {"model": {"provider": "p"}}
app_model = _app()
app_model_config = _config()
annotation_reply = {"enabled": False}
conversation = mocker.MagicMock()
conversation = _conversation()
to_dict = mocker.patch.object(
AppModelConfig,
"to_dict",
return_value={"model": {"provider": "p"}},
)
mocker.patch("core.app.apps.agent_chat.app_config_manager.ModelConfigManager.convert")
mocker.patch("core.app.apps.agent_chat.app_config_manager.PromptTemplateConfigManager.convert")
@@ -78,15 +126,19 @@ class TestAgentChatAppConfigManagerGetAppConfig:
annotation_reply=annotation_reply,
)
assert result.app_model_config_dict == app_model_config.to_dict.return_value
assert result.app_model_config_dict == {"model": {"provider": "p"}}
assert result.app_model_config_from.value == "conversation-specific-config"
app_model_config.to_dict.assert_called_once_with(annotation_reply=annotation_reply)
to_dict.assert_called_once_with(annotation_reply=annotation_reply)
def test_get_app_config_latest_config(self, mocker: MockerFixture):
app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent-chat")
app_model_config = mocker.MagicMock(id="cfg1")
app_model_config.to_dict.return_value = {"model": {"provider": "p"}}
app_model = _app()
app_model_config = _config()
annotation_reply = {"enabled": False}
to_dict = mocker.patch.object(
AppModelConfig,
"to_dict",
return_value={"model": {"provider": "p"}},
)
mocker.patch("core.app.apps.agent_chat.app_config_manager.ModelConfigManager.convert")
mocker.patch("core.app.apps.agent_chat.app_config_manager.PromptTemplateConfigManager.convert")
@@ -112,7 +164,7 @@ class TestAgentChatAppConfigManagerGetAppConfig:
)
assert result.app_model_config_from.value == "app-latest-config"
app_model_config.to_dict.assert_called_once_with(annotation_reply=annotation_reply)
to_dict.assert_called_once_with(annotation_reply=annotation_reply)
class TestAgentChatAppConfigManagerConfigValidate:
@@ -186,7 +238,7 @@ class TestAgentChatAppConfigManagerConfigValidate:
side_effect=lambda tenant_id, cfg: return_with_key("moderation"),
)
filtered = AgentChatAppConfigManager.config_validate("tenant", config, mocker.MagicMock())
filtered = AgentChatAppConfigManager.config_validate("tenant", config, _session())
assert set(filtered.keys()) == {
"model",
"user_input_form",
@@ -207,7 +259,7 @@ class TestAgentChatAppConfigManagerConfigValidate:
class TestValidateAgentModeAndSetDefaults:
def test_defaults_when_missing(self):
config = {}
updated, keys = AgentChatAppConfigManager.validate_agent_mode_and_set_defaults("tenant", config, MagicMock())
updated, keys = AgentChatAppConfigManager.validate_agent_mode_and_set_defaults("tenant", config, _session())
assert "agent_mode" in updated
assert updated["agent_mode"]["enabled"] is False
assert updated["agent_mode"]["tools"] == []
@@ -220,37 +272,37 @@ class TestValidateAgentModeAndSetDefaults:
def test_agent_mode_type_validation(self, agent_mode):
with pytest.raises(ValueError):
AgentChatAppConfigManager.validate_agent_mode_and_set_defaults(
"tenant", {"agent_mode": agent_mode}, MagicMock()
"tenant", {"agent_mode": agent_mode}, _session()
)
def test_agent_mode_empty_list_defaults(self):
config = {"agent_mode": []}
updated, _ = AgentChatAppConfigManager.validate_agent_mode_and_set_defaults("tenant", config, MagicMock())
updated, _ = AgentChatAppConfigManager.validate_agent_mode_and_set_defaults("tenant", config, _session())
assert updated["agent_mode"]["enabled"] is False
assert updated["agent_mode"]["tools"] == []
def test_enabled_must_be_bool(self):
with pytest.raises(ValueError):
AgentChatAppConfigManager.validate_agent_mode_and_set_defaults(
"tenant", {"agent_mode": {"enabled": "yes"}}, MagicMock()
"tenant", {"agent_mode": {"enabled": "yes"}}, _session()
)
def test_strategy_must_be_valid(self):
with pytest.raises(ValueError):
AgentChatAppConfigManager.validate_agent_mode_and_set_defaults(
"tenant", {"agent_mode": {"enabled": True, "strategy": "invalid"}}, MagicMock()
"tenant", {"agent_mode": {"enabled": True, "strategy": "invalid"}}, _session()
)
def test_tools_must_be_list(self):
with pytest.raises(ValueError):
AgentChatAppConfigManager.validate_agent_mode_and_set_defaults(
"tenant", {"agent_mode": {"enabled": True, "tools": "not-list"}}, MagicMock()
"tenant", {"agent_mode": {"enabled": True, "tools": "not-list"}}, _session()
)
def test_old_tool_dataset_requires_id(self):
with pytest.raises(ValueError):
AgentChatAppConfigManager.validate_agent_mode_and_set_defaults(
"tenant", {"agent_mode": {"enabled": True, "tools": [{"dataset": {"enabled": True}}]}}, MagicMock()
"tenant", {"agent_mode": {"enabled": True, "tools": [{"dataset": {"enabled": True}}]}}, _session()
)
def test_old_tool_dataset_id_must_be_uuid(self):
@@ -258,7 +310,7 @@ class TestValidateAgentModeAndSetDefaults:
AgentChatAppConfigManager.validate_agent_mode_and_set_defaults(
"tenant",
{"agent_mode": {"enabled": True, "tools": [{"dataset": {"enabled": True, "id": "bad"}}]}},
MagicMock(),
_session(),
)
def test_old_tool_dataset_id_not_exists(self, mocker: MockerFixture):
@@ -271,7 +323,7 @@ class TestValidateAgentModeAndSetDefaults:
AgentChatAppConfigManager.validate_agent_mode_and_set_defaults(
"tenant",
{"agent_mode": {"enabled": True, "tools": [{"dataset": {"enabled": True, "id": dataset_id}}]}},
MagicMock(),
_session(),
)
def test_old_tool_enabled_must_be_bool(self):
@@ -279,7 +331,7 @@ class TestValidateAgentModeAndSetDefaults:
AgentChatAppConfigManager.validate_agent_mode_and_set_defaults(
"tenant",
{"agent_mode": {"enabled": True, "tools": [{"dataset": {"enabled": "yes", "id": str(uuid.uuid4())}}]}},
MagicMock(),
_session(),
)
@pytest.mark.parametrize("missing_key", ["provider_type", "provider_id", "tool_name", "tool_parameters"])
@@ -288,7 +340,7 @@ class TestValidateAgentModeAndSetDefaults:
tool.pop(missing_key, None)
with pytest.raises(ValueError):
AgentChatAppConfigManager.validate_agent_mode_and_set_defaults(
"tenant", {"agent_mode": {"enabled": True, "tools": [tool]}}, MagicMock()
"tenant", {"agent_mode": {"enabled": True, "tools": [tool]}}, _session()
)
def test_valid_old_and_new_style_tools(self, mocker: MockerFixture):
@@ -313,6 +365,6 @@ class TestValidateAgentModeAndSetDefaults:
}
}
updated, _ = AgentChatAppConfigManager.validate_agent_mode_and_set_defaults("tenant", config, MagicMock())
updated, _ = AgentChatAppConfigManager.validate_agent_mode_and_set_defaults("tenant", config, _session())
assert updated["agent_mode"]["tools"][0]["dataset"]["enabled"] is False
assert updated["agent_mode"]["tools"][1]["enabled"] is False
@@ -1,23 +1,83 @@
import contextlib
import inspect
import logging
from decimal import Decimal
import pytest
from pydantic import ValidationError
from pytest_mock import MockerFixture
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.orm import Session
from core.app.apps.agent_chat.app_generator import AgentChatAppGenerator
from core.app.apps.exc import GenerateTaskStoppedError
from core.app.entities.app_invoke_entities import InvokeFrom
from graphon.model_runtime.errors.invoke import InvokeAuthorizationError
from models import Account
from models.enums import ConversationFromSource
from models.model import App, AppMode, AppModelConfig, Conversation, Message
from tests.unit_tests.config_override import apply_config_overrides
class DummyAccount:
def __init__(self, user_id):
self.id = user_id
self.session_id = f"session-{user_id}"
def _app() -> App:
return App(
id="app1",
tenant_id="tenant",
name="Agent chat app",
description="",
mode=AppMode.AGENT_CHAT,
enable_site=False,
enable_api=False,
)
def _account() -> Account:
account = Account(name="User", email="user@example.com")
account.id = "user"
return account
def _conversation() -> Conversation:
conversation = Conversation(
id="conv",
app_id="app1",
app_model_config_id=None,
model_provider=None,
override_model_configs=None,
model_id=None,
mode=AppMode.AGENT_CHAT,
name="Conversation",
inputs={},
introduction="",
system_instruction="",
system_instruction_tokens=0,
status="normal",
invoke_from=InvokeFrom.WEB_APP,
from_source=ConversationFromSource.CONSOLE,
from_end_user_id=None,
from_account_id="user",
)
return conversation
def _message() -> Message:
return Message(
id="msg",
app_id="app1",
conversation_id="conv",
inputs={},
query="hello",
message={},
message_unit_price=Decimal(0),
answer="",
answer_unit_price=Decimal(0),
total_price=Decimal(0),
currency="USD",
invoke_from=InvokeFrom.WEB_APP,
from_source=ConversationFromSource.CONSOLE,
from_end_user_id=None,
from_account_id="user",
app_mode=AppMode.AGENT_CHAT,
)
@pytest.fixture
@@ -33,8 +93,8 @@ def generator(mocker: MockerFixture):
class TestAgentChatAppGeneratorGenerate:
def test_generate_rejects_blocking_mode(self, generator, mocker: MockerFixture, sqlite_session: Session):
app_model = mocker.MagicMock()
user = DummyAccount("user")
app_model = _app()
user = _account()
with pytest.raises(ValueError):
generator.generate(
session=sqlite_session,
@@ -46,8 +106,8 @@ class TestAgentChatAppGeneratorGenerate:
)
def test_generate_requires_query(self, generator, mocker: MockerFixture, sqlite_session: Session):
app_model = mocker.MagicMock()
user = DummyAccount("user")
app_model = _app()
user = _account()
with pytest.raises(ValueError):
generator.generate(
session=sqlite_session,
@@ -58,8 +118,8 @@ class TestAgentChatAppGeneratorGenerate:
)
def test_generate_rejects_non_string_query(self, generator, mocker: MockerFixture, sqlite_session: Session):
app_model = mocker.MagicMock()
user = DummyAccount("user")
app_model = _app()
user = _account()
with pytest.raises(ValueError):
generator.generate(
session=sqlite_session,
@@ -70,9 +130,9 @@ class TestAgentChatAppGeneratorGenerate:
)
def test_generate_override_requires_debugger(self, generator, mocker: MockerFixture, sqlite_session: Session):
app_model = mocker.MagicMock()
user = DummyAccount("user")
generator._get_app_model_config = mocker.MagicMock(return_value=mocker.MagicMock())
app_model = _app()
user = _account()
generator._get_app_model_config = mocker.MagicMock(return_value=AppModelConfig(app_id="app1"))
with pytest.raises(ValueError):
generator.generate(
@@ -84,18 +144,15 @@ class TestAgentChatAppGeneratorGenerate:
)
def test_generate_success_with_debugger_override(self, generator, mocker: MockerFixture, sqlite_session: Session):
app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent-chat")
app_model_config = mocker.MagicMock(id="cfg1")
app_model_config.to_dict.return_value = {"model": {"provider": "p"}}
app_model = _app()
app_model_config = AppModelConfig(app_id="app1")
user = DummyAccount("user")
user = _account()
invoke_from = InvokeFrom.DEBUGGER
generator._get_app_model_config = mocker.MagicMock(return_value=app_model_config)
generator._prepare_user_inputs = mocker.MagicMock(return_value={"x": 1})
generator._init_generate_records = mocker.MagicMock(
return_value=(mocker.MagicMock(id="conv", mode="agent-chat"), mocker.MagicMock(id="msg"))
)
generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message()))
generator._handle_response = mocker.MagicMock(return_value="response")
mocker.patch(
@@ -121,7 +178,7 @@ class TestAgentChatAppGeneratorGenerate:
)
mocker.patch(
"core.app.apps.agent_chat.app_generator.ConversationService.get_conversation",
return_value=mocker.MagicMock(id="conv"),
return_value=_conversation(),
)
mocker.patch(
"core.app.apps.agent_chat.app_generator.TraceQueueManager",
@@ -177,20 +234,19 @@ class TestAgentChatAppGeneratorGenerate:
thread_obj.start.assert_called_once()
def test_generate_without_file_config(self, generator, mocker: MockerFixture, sqlite_session: Session):
app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent-chat")
app_model_config = mocker.MagicMock(id="cfg1", app_id="app1")
app_model_config.to_dict.return_value = {"model": {"provider": "p"}}
app_model = _app()
app_model_config = AppModelConfig(app_id="app1")
annotation_reply = {"enabled": False}
user = DummyAccount("user")
user = _account()
generator._get_app_model_config = mocker.MagicMock(return_value=app_model_config)
generator._prepare_user_inputs = mocker.MagicMock(return_value={"x": 1})
generator._init_generate_records = mocker.MagicMock(
return_value=(mocker.MagicMock(id="conv", mode="agent-chat"), mocker.MagicMock(id="msg"))
)
generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message()))
generator._handle_response = mocker.MagicMock(return_value="response")
to_dict = mocker.patch.object(AppModelConfig, "to_dict", return_value={"model": {"provider": "p"}})
load_annotation_reply_config = mocker.patch(
"core.app.apps.agent_chat.app_generator.load_annotation_reply_config",
return_value=annotation_reply,
@@ -251,7 +307,7 @@ class TestAgentChatAppGeneratorGenerate:
assert result == {"result": "ok"}
load_annotation_reply_config.assert_called_once_with(session, "app1")
app_model_config.to_dict.assert_called_once_with(annotation_reply=annotation_reply)
to_dict.assert_called_once_with(annotation_reply=annotation_reply)
assert get_app_config.call_args.kwargs["annotation_reply"] is annotation_reply
@@ -264,20 +320,14 @@ class TestAgentChatAppGeneratorWorker:
mocker.patch("core.app.apps.agent_chat.app_generator.preserve_flask_contexts", ctx_manager)
def test_generate_worker_handles_generate_task_stopped(
self, generator, mocker: MockerFixture, sqlite_session_factory: sessionmaker[Session]
):
def test_generate_worker_handles_generate_task_stopped(self, generator, mocker: MockerFixture):
queue_manager = mocker.MagicMock()
generator._get_conversation = mocker.MagicMock(return_value=mocker.MagicMock())
generator._get_message = mocker.MagicMock(return_value=mocker.MagicMock())
generator._get_conversation = mocker.MagicMock(return_value=_conversation())
generator._get_message = mocker.MagicMock(return_value=_message())
runner = mocker.MagicMock()
runner.run.side_effect = GenerateTaskStoppedError()
mocker.patch("core.app.apps.agent_chat.app_generator.AgentChatAppRunner", return_value=runner)
mocker.patch(
"core.app.apps.agent_chat.app_generator.session_factory.create_session",
side_effect=sqlite_session_factory,
)
generator._generate_worker(
flask_app=mocker.MagicMock(),
@@ -299,20 +349,14 @@ class TestAgentChatAppGeneratorWorker:
Exception("bad"),
],
)
def test_generate_worker_publishes_errors(
self, generator, mocker: MockerFixture, error, sqlite_session_factory: sessionmaker[Session]
):
def test_generate_worker_publishes_errors(self, generator, mocker: MockerFixture, error):
queue_manager = mocker.MagicMock()
generator._get_conversation = mocker.MagicMock(return_value=mocker.MagicMock())
generator._get_message = mocker.MagicMock(return_value=mocker.MagicMock())
generator._get_conversation = mocker.MagicMock(return_value=_conversation())
generator._get_message = mocker.MagicMock(return_value=_message())
runner = mocker.MagicMock()
runner.run.side_effect = error
mocker.patch("core.app.apps.agent_chat.app_generator.AgentChatAppRunner", return_value=runner)
mocker.patch(
"core.app.apps.agent_chat.app_generator.session_factory.create_session",
side_effect=sqlite_session_factory,
)
generator._generate_worker(
flask_app=mocker.MagicMock(),
@@ -331,19 +375,14 @@ class TestAgentChatAppGeneratorWorker:
mocker: MockerFixture,
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
sqlite_session_factory: sessionmaker[Session],
):
queue_manager = mocker.MagicMock()
generator._get_conversation = mocker.MagicMock(return_value=mocker.MagicMock())
generator._get_message = mocker.MagicMock(return_value=mocker.MagicMock())
generator._get_conversation = mocker.MagicMock(return_value=_conversation())
generator._get_message = mocker.MagicMock(return_value=_message())
runner = mocker.MagicMock()
runner.run.side_effect = ValueError("bad")
mocker.patch("core.app.apps.agent_chat.app_generator.AgentChatAppRunner", return_value=runner)
mocker.patch(
"core.app.apps.agent_chat.app_generator.session_factory.create_session",
side_effect=sqlite_session_factory,
)
apply_config_overrides(monkeypatch, DEBUG=True)
@@ -1,17 +1,21 @@
from decimal import Decimal
from types import SimpleNamespace
from unittest.mock import ANY, MagicMock
import pytest
from pytest_mock import MockerFixture
from sqlalchemy import event
from sqlalchemy.orm import Session
import core.app.apps.completion.app_runner as module
from core.app.apps.completion.app_runner import CompletionAppRunner
from core.app.entities.app_invoke_entities import InvokeFrom
from core.credit_usage import CreditUsageAppType, CreditUsageCreatedBy
from core.moderation.base import ModerationError
from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent
from graphon.model_runtime.entities.model_entities import ModelType
from models.model import App, AppMode, IconType
from models.enums import ConversationFromSource
from models.model import App, AppMode, Conversation, IconType, Message
APP_ID = "00000000-0000-0000-0000-000000000001"
TENANT_ID = "00000000-0000-0000-0000-000000000002"
@@ -56,7 +60,27 @@ def _build_generate_entity(app_config, file_upload_config=None):
)
def _persist_app(session: Session) -> App:
def _message() -> Message:
return Message(
id="msg",
app_id=APP_ID,
conversation_id="conv",
inputs={},
query="query",
message={},
message_unit_price=Decimal(0),
answer="",
answer_unit_price=Decimal(0),
total_price=Decimal(0),
currency="USD",
invoke_from=InvokeFrom.SERVICE_API,
from_source=ConversationFromSource.API,
from_account_id="user",
app_mode=AppMode.COMPLETION,
)
def _persist_records(session: Session) -> tuple[App, Message]:
app = App(
id=APP_ID,
tenant_id=TENANT_ID,
@@ -68,9 +92,20 @@ def _persist_app(session: Session) -> App:
enable_site=False,
enable_api=False,
)
session.add(app)
conversation = Conversation(
id="conv",
app_id=APP_ID,
mode=AppMode.COMPLETION,
name="Conversation",
inputs={},
invoke_from=InvokeFrom.SERVICE_API,
from_source=ConversationFromSource.API,
from_account_id="user",
)
message = _message()
session.add_all([app, conversation, message])
session.commit()
return app
return app, message
class TestCompletionAppRunner:
@@ -79,10 +114,10 @@ class TestCompletionAppRunner:
app_generate_entity = _build_generate_entity(app_config)
with pytest.raises(ValueError):
runner.run(app_generate_entity, MagicMock(), MagicMock(), sqlite_session)
runner.run(app_generate_entity, MagicMock(), _message(), sqlite_session)
def test_run_moderation_error_outputs_direct(self, runner, mocker: MockerFixture, sqlite_session: Session):
_persist_app(sqlite_session)
_, message = _persist_records(sqlite_session)
app_config = _build_app_config()
app_generate_entity = _build_generate_entity(app_config)
@@ -92,13 +127,13 @@ class TestCompletionAppRunner:
runner.direct_output = MagicMock()
runner._handle_invoke_result = MagicMock()
runner.run(app_generate_entity, MagicMock(), MagicMock(id="msg"), sqlite_session)
runner.run(app_generate_entity, MagicMock(), message, sqlite_session)
runner.direct_output.assert_called_once()
runner._handle_invoke_result.assert_not_called()
def test_run_hosting_moderation_stops(self, runner, mocker: MockerFixture, sqlite_session: Session):
_persist_app(sqlite_session)
_, message = _persist_records(sqlite_session)
app_config = _build_app_config()
app_generate_entity = _build_generate_entity(app_config)
@@ -108,12 +143,12 @@ class TestCompletionAppRunner:
runner.check_hosting_moderation = MagicMock(return_value=True)
runner._handle_invoke_result = MagicMock()
runner.run(app_generate_entity, MagicMock(), MagicMock(id="msg"), sqlite_session)
runner.run(app_generate_entity, MagicMock(), message, sqlite_session)
runner._handle_invoke_result.assert_not_called()
def test_run_dataset_and_external_tools_flow(self, runner, mocker: MockerFixture, sqlite_session: Session):
_persist_app(sqlite_session)
_, message = _persist_records(sqlite_session)
retrieve_config = MagicMock(query_variable="qvar")
dataset_config = MagicMock(dataset_ids=["ds"], retrieve_config=retrieve_config)
@@ -146,7 +181,7 @@ class TestCompletionAppRunner:
model_manager.get_model_instance.return_value = model_instance
mocker.patch.object(module.ModelManager, "for_tenant", return_value=model_manager)
runner.run(app_generate_entity, MagicMock(), MagicMock(id="msg", tenant_id=TENANT_ID), sqlite_session)
runner.run(app_generate_entity, MagicMock(), message, sqlite_session)
dataset_retrieval.retrieve.assert_called_once()
assert dataset_retrieval.retrieve.call_args.kwargs["query"] == "query_from_input"
@@ -155,26 +190,22 @@ class TestCompletionAppRunner:
def test_run_closes_explicit_session_before_stream_consumption(
self, runner, mocker: MockerFixture, sqlite_session: Session
):
_persist_app(sqlite_session)
_, message = _persist_records(sqlite_session)
app_config = _build_app_config()
app_generate_entity = _build_generate_entity(app_config)
queue_manager = MagicMock()
events = []
session = sqlite_session
original_commit = session.commit
original_close = session.close
def commit_session() -> None:
def record_commit(_session: Session) -> None:
events.append("commit")
original_commit()
def close_session() -> None:
def record_close() -> None:
events.append("close")
original_close()
mocker.patch.object(session, "commit", side_effect=commit_session)
mocker.patch.object(session, "close", side_effect=close_session)
runner.organize_prompt_messages = MagicMock(return_value=([], None))
runner.moderation_for_inputs = MagicMock(return_value=(None, app_generate_entity.inputs, "query"))
runner.check_hosting_moderation = MagicMock(return_value=False)
@@ -195,8 +226,13 @@ class TestCompletionAppRunner:
model_manager = MagicMock()
model_manager.get_model_instance.return_value = model_instance
mocker.patch.object(module.ModelManager, "for_tenant", return_value=model_manager)
mocker.patch.object(session, "close", side_effect=record_close)
runner.run(app_generate_entity, queue_manager, MagicMock(id="msg"), session)
event.listen(session, "after_commit", record_commit)
try:
runner.run(app_generate_entity, queue_manager, message, session)
finally:
event.remove(session, "after_commit", record_commit)
assert events == ["commit", "close", "invoke", "first-chunk"]
runner._handle_invoke_result.assert_called_once_with(
@@ -216,7 +252,7 @@ class TestCompletionAppRunner:
sqlite_session: Session,
stream: bool,
):
_persist_app(sqlite_session)
_, message = _persist_records(sqlite_session)
app_config = _build_app_config()
app_generate_entity = _build_generate_entity(app_config)
app_generate_entity.stream = stream
@@ -233,7 +269,7 @@ class TestCompletionAppRunner:
model_manager.get_model_instance.return_value = model_instance
model_manager_factory = mocker.patch.object(module.ModelManager, "for_tenant", return_value=model_manager)
runner.run(app_generate_entity, MagicMock(), MagicMock(id="msg"), sqlite_session)
runner.run(app_generate_entity, MagicMock(), message, sqlite_session)
model_manager_factory.assert_called_once_with(tenant_id=TENANT_ID)
model_manager.get_model_instance.assert_called_once_with(
@@ -255,7 +291,7 @@ class TestCompletionAppRunner:
)
def test_run_uses_low_image_detail_default(self, runner, mocker: MockerFixture, sqlite_session: Session):
_persist_app(sqlite_session)
_, message = _persist_records(sqlite_session)
app_config = _build_app_config()
app_generate_entity = _build_generate_entity(app_config, file_upload_config=None)
@@ -264,7 +300,7 @@ class TestCompletionAppRunner:
runner.moderation_for_inputs = MagicMock(return_value=(None, app_generate_entity.inputs, "query"))
runner.check_hosting_moderation = MagicMock(return_value=True)
runner.run(app_generate_entity, MagicMock(), MagicMock(id="msg"), sqlite_session)
runner.run(app_generate_entity, MagicMock(), message, sqlite_session)
assert (
runner.organize_prompt_messages.call_args.kwargs["image_detail_config"]
@@ -1,19 +1,34 @@
from types import SimpleNamespace
from unittest.mock import MagicMock
from pytest_mock import MockerFixture
from sqlalchemy.orm import Session
import core.app.apps.completion.app_config_manager as module
from core.app.app_config.entities import EasyUIBasedAppModelConfigFrom
from core.app.apps.completion.app_config_manager import CompletionAppConfigManager
from models.model import AppMode
from models.model import App, AppMode, AppModelConfig
def _app() -> App:
return App(
id="app1",
tenant_id="tenant",
name="Completion app",
description="",
mode=AppMode.COMPLETION,
enable_site=False,
enable_api=False,
)
def _config() -> AppModelConfig:
return AppModelConfig(app_id="app1")
class TestCompletionAppConfigManager:
def test_get_app_config_with_override(self, mocker: MockerFixture):
app_model = MagicMock(tenant_id="tenant", id="app1", mode=AppMode.COMPLETION.value)
app_model_config = MagicMock(id="cfg1")
app_model_config.to_dict.return_value = {"model": {"provider": "x"}}
app_model = _app()
app_model_config = _config()
override_config = {"model": {"provider": "override"}}
@@ -39,10 +54,14 @@ class TestCompletionAppConfigManager:
assert result.app_mode == AppMode.COMPLETION
def test_get_app_config_without_override_uses_model_config(self, mocker: MockerFixture):
app_model = MagicMock(tenant_id="tenant", id="app1", mode=AppMode.COMPLETION)
app_model_config = MagicMock(id="cfg1")
app_model_config.to_dict.return_value = {"model": {"provider": "x"}}
app_model = _app()
app_model_config = _config()
annotation_reply = {"enabled": False}
to_dict = mocker.patch.object(
AppModelConfig,
"to_dict",
return_value={"model": {"provider": "x"}},
)
mocker.patch.object(module.ModelConfigManager, "convert", return_value="model")
mocker.patch.object(module.PromptTemplateConfigManager, "convert", return_value="prompt")
@@ -60,9 +79,9 @@ class TestCompletionAppConfigManager:
assert result.app_model_config_from == EasyUIBasedAppModelConfigFrom.APP_LATEST_CONFIG
assert result.app_model_config_dict == {"model": {"provider": "x"}}
app_model_config.to_dict.assert_called_once_with(annotation_reply=annotation_reply)
to_dict.assert_called_once_with(annotation_reply=annotation_reply)
def test_config_validate_filters_related_keys(self, mocker: MockerFixture):
def test_config_validate_filters_related_keys(self, mocker: MockerFixture, unbound_session: Session):
config = {
"model": {"provider": "x"},
"variables": ["v"],
@@ -116,7 +135,7 @@ class TestCompletionAppConfigManager:
return_value=(config, ["moderation"]),
)
filtered = CompletionAppConfigManager.config_validate("tenant", config, MagicMock())
filtered = CompletionAppConfigManager.config_validate("tenant", config, unbound_session)
assert "extra" not in filtered
assert set(filtered.keys()) == {
@@ -13,8 +13,9 @@ from core.app.apps.completion.app_generator import CompletionAppGenerator
from core.app.apps.exc import GenerateTaskStoppedError
from core.app.entities.app_invoke_entities import InvokeFrom
from graphon.model_runtime.errors.invoke import InvokeAuthorizationError
from models import Account
from models.enums import ConversationFromSource
from models.model import AppMode, AppModelConfig, Conversation, Message
from models.model import App, AppMode, AppModelConfig, Conversation, Message
from services.errors.app import MoreLikeThisDisabledError
from services.errors.message import MessageNotExistsError
@@ -44,18 +45,32 @@ def generator(mocker: MockerFixture):
return gen
def _build_app_model():
return MagicMock(tenant_id="tenant", id="app1", mode="completion", app_model_config_id="cfg-current")
@pytest.fixture(autouse=True)
def _bind_db_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(module.db, "session", sqlite_session)
def _build_user():
return MagicMock(id="user", session_id="session")
def _build_app_model() -> App:
return App(
id="app1",
tenant_id="tenant",
name="Completion app",
description="",
mode=AppMode.COMPLETION,
app_model_config_id="cfg-current",
enable_site=False,
enable_api=False,
)
def _build_app_model_config():
config = MagicMock(id="cfg", app_id="app1")
config.to_dict.return_value = {"model": {"provider": "x"}}
return config
def _build_user() -> Account:
account = Account(name="User", email="user@example.com")
account.id = "user"
return account
def _build_app_model_config() -> AppModelConfig:
return AppModelConfig(app_id="app1", model='{"provider": "x"}')
def _persist_message(
@@ -80,7 +95,7 @@ def _persist_message(
invoke_from=InvokeFrom.WEB_APP,
from_source=ConversationFromSource.CONSOLE,
from_end_user_id=None,
from_account_id=None,
from_account_id="user",
read_at=None,
read_account_id=None,
)
@@ -108,7 +123,7 @@ def _persist_message(
invoke_from=InvokeFrom.WEB_APP,
from_source=ConversationFromSource.CONSOLE,
from_end_user_id=None,
from_account_id=None,
from_account_id="user",
workflow_run_id=None,
app_mode=AppMode.COMPLETION,
)
@@ -144,6 +159,7 @@ class TestCompletionAppGenerator:
def test_generate_success_no_file_config(self, generator, mocker: MockerFixture, sqlite_session: Session):
app_model_config = _build_app_model_config()
mocker.patch.object(generator, "_get_app_model_config", return_value=app_model_config)
to_dict = mocker.patch.object(AppModelConfig, "to_dict", return_value={"model": {"provider": "x"}})
annotation_reply = {"enabled": False}
load_annotation_reply_config = mocker.patch.object(
module,
@@ -163,8 +179,27 @@ class TestCompletionAppGenerator:
mocker.patch.object(generator, "_prepare_user_inputs", return_value={"k": "v"})
conversation = MagicMock(id="conv", mode="completion")
message = MagicMock(id="msg")
conversation = Conversation(
id="conv",
app_id="app1",
mode=AppMode.COMPLETION,
name="Conversation",
inputs={},
from_source=ConversationFromSource.CONSOLE,
)
message = Message(
id="msg",
app_id="app1",
conversation_id="conv",
inputs={},
query="q",
message={},
message_unit_price=Decimal(0),
answer="",
answer_unit_price=Decimal(0),
currency="USD",
from_source=ConversationFromSource.CONSOLE,
)
mocker.patch.object(generator, "_init_generate_records", return_value=(conversation, message))
mocker.patch.object(generator, "_handle_response", return_value="response")
@@ -183,12 +218,13 @@ class TestCompletionAppGenerator:
assert generator.generate_entity.call_args.kwargs["extras"]["trace_session_id"] == "session-1"
module.file_factory.build_from_mappings.assert_not_called()
load_annotation_reply_config.assert_called_once_with(sqlite_session, "app1")
app_model_config.to_dict.assert_called_once_with(annotation_reply=annotation_reply)
to_dict.assert_called_once_with(annotation_reply=annotation_reply)
assert get_app_config.call_args.kwargs["annotation_reply"] is annotation_reply
def test_generate_success_with_files(self, generator, mocker: MockerFixture, sqlite_session: Session):
app_model_config = _build_app_model_config()
mocker.patch.object(generator, "_get_app_model_config", return_value=app_model_config)
mocker.patch.object(AppModelConfig, "to_dict", return_value={"model": {"provider": "x"}})
file_extra_config = MagicMock()
mocker.patch.object(module.FileUploadConfigManager, "convert", return_value=file_extra_config)
@@ -200,8 +236,27 @@ class TestCompletionAppGenerator:
mocker.patch.object(generator, "_prepare_user_inputs", return_value={"k": "v"})
conversation = MagicMock(id="conv", mode="completion")
message = MagicMock(id="msg")
conversation = Conversation(
id="conv",
app_id="app1",
mode=AppMode.COMPLETION,
name="Conversation",
inputs={},
from_source=ConversationFromSource.CONSOLE,
)
message = Message(
id="msg",
app_id="app1",
conversation_id="conv",
inputs={},
query="q",
message={},
message_unit_price=Decimal(0),
answer="",
answer_unit_price=Decimal(0),
currency="USD",
from_source=ConversationFromSource.CONSOLE,
)
mocker.patch.object(generator, "_init_generate_records", return_value=(conversation, message))
mocker.patch.object(generator, "_handle_response", return_value="response")
@@ -238,7 +293,29 @@ class TestCompletionAppGenerator:
mocker.patch.object(
generator,
"_init_generate_records",
return_value=(MagicMock(id="conv", mode="completion"), MagicMock(id="msg")),
return_value=(
Conversation(
id="conv",
app_id="app1",
mode=AppMode.COMPLETION,
name="Conversation",
inputs={},
from_source=ConversationFromSource.CONSOLE,
),
Message(
id="msg",
app_id="app1",
conversation_id="conv",
inputs={},
query="q",
message={},
message_unit_price=Decimal(0),
answer="",
answer_unit_price=Decimal(0),
currency="USD",
from_source=ConversationFromSource.CONSOLE,
),
),
)
mocker.patch.object(generator, "_handle_response", return_value="response")
mocker.patch.object(module.CompletionAppGenerateResponseConverter, "convert", return_value="converted")
@@ -350,7 +427,29 @@ class TestCompletionAppGenerator:
mocker.patch.object(
generator,
"_init_generate_records",
return_value=(MagicMock(id="conv", mode="completion"), MagicMock(id="msg")),
return_value=(
Conversation(
id="conv",
app_id="app1",
mode=AppMode.COMPLETION,
name="Conversation",
inputs={},
from_source=ConversationFromSource.CONSOLE,
),
Message(
id="msg",
app_id="app1",
conversation_id="conv",
inputs={},
query="q",
message={},
message_unit_price=Decimal(0),
answer="",
answer_unit_price=Decimal(0),
currency="USD",
from_source=ConversationFromSource.CONSOLE,
),
),
)
mocker.patch.object(generator, "_handle_response", return_value="response")
mocker.patch.object(module.CompletionAppGenerateResponseConverter, "convert", return_value="converted")
@@ -399,13 +498,8 @@ class TestCompletionAppGenerator:
flask_app = MagicMock()
flask_app.app_context.return_value = contextlib.nullcontext()
session_context = mocker.MagicMock()
session_context.__enter__.return_value = sqlite_session
create_session = mocker.patch.object(module.session_factory, "create_session")
create_session.return_value = session_context
mocker.patch.object(module.db, "session")
mocker.patch.object(generator, "_get_message", return_value=MagicMock())
message = _persist_message(sqlite_session)
mocker.patch.object(generator, "_get_message", return_value=message)
runner_instance = MagicMock()
runner_instance.run.side_effect = error
@@ -5,16 +5,18 @@ from unittest.mock import MagicMock
import pytest
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, scoped_session, sessionmaker
from sqlalchemy.orm import Session
import core.app.apps.advanced_chat.app_generator as app_generator_module
from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig
from core.app.apps.advanced_chat.app_generator import AdvancedChatAppGenerator
from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom
from core.app.task_pipeline import message_cycle_manager
from core.app.task_pipeline.message_cycle_manager import MessageCycleManager
from core.ops.ops_trace_manager import TraceQueueManager
from models import Account, Workflow
from models.enums import ConversationFromSource
from models.model import AppMode, Conversation
from models.model import App, AppMode, Conversation
from services.errors.conversation import ConversationNotExistsError
@@ -47,6 +49,39 @@ def _make_generate_entity(app_config: WorkflowUIBasedAppConfig) -> AdvancedChatA
)
def _app() -> App:
return App(
id="app-id",
tenant_id="tenant-id",
name="Advanced chat app",
description="",
mode=AppMode.ADVANCED_CHAT,
enable_site=False,
enable_api=False,
)
def _workflow() -> Workflow:
return Workflow.new(
tenant_id="tenant-id",
app_id="app-id",
type="chat",
version=Workflow.VERSION_DRAFT,
graph='{"nodes": [], "edges": []}',
features="{}",
created_by="user-id",
environment_variables=[],
conversation_variables=[],
rag_pipeline_variables=[],
)
def _account() -> Account:
account = Account(name="User", email="user@example.com")
account.id = "user-id"
return account
def test_init_generate_records_sets_conversation_metadata(sqlite_session: Session):
app_config = _make_app_config()
entity = _make_generate_entity(app_config)
@@ -107,17 +142,11 @@ def test_generate_falls_back_to_new_conversation_when_conversation_missing(
monkeypatch: pytest.MonkeyPatch,
sqlite_engine: Engine,
sqlite_session: Session,
sqlite_session_factory: sessionmaker[Session],
):
app_config = _make_app_config()
workflow = SimpleNamespace(
features_dict={},
tenant_id="tenant-id",
app_id="app-id",
id="workflow-id",
)
app_model = SimpleNamespace(id="app-id", tenant_id="tenant-id")
user = SimpleNamespace(id="user-id", session_id="session-id")
workflow = _workflow()
app_model = _app()
user = _account()
def raise_conversation_not_exists(**_kwargs):
raise ConversationNotExistsError()
@@ -134,11 +163,7 @@ def test_generate_falls_back_to_new_conversation_when_conversation_missing(
"core.app.apps.advanced_chat.app_generator.AdvancedChatAppConfigManager.get_app_config",
lambda **_kwargs: app_config,
)
db_session = scoped_session(sqlite_session_factory)
monkeypatch.setattr(
"core.app.apps.advanced_chat.app_generator.db",
SimpleNamespace(engine=sqlite_engine, session=db_session),
)
monkeypatch.setattr(app_generator_module, "db", SimpleNamespace(engine=sqlite_engine))
trace_manager = object.__new__(TraceQueueManager)
monkeypatch.setattr(
"core.app.apps.advanced_chat.app_generator.TraceQueueManager",
@@ -179,7 +204,6 @@ def test_generate_falls_back_to_new_conversation_when_conversation_missing(
application_generate_entity = captured["application_generate_entity"]
assert isinstance(application_generate_entity, AdvancedChatAppGenerateEntity)
assert application_generate_entity.conversation_id is None
db_session.remove()
def test_message_cycle_manager_uses_new_conversation_flag(monkeypatch: pytest.MonkeyPatch):
@@ -103,19 +103,18 @@ def _runtime_state() -> GraphRuntimeState:
def _repositories(
sqlite_session: Session, app: App, end_user: EndUser
sqlite_session_factory: sessionmaker[Session], app: App, end_user: EndUser
) -> tuple[SQLAlchemyWorkflowExecutionRepository, SQLAlchemyWorkflowNodeExecutionRepository]:
session_factory = sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False)
return (
SQLAlchemyWorkflowExecutionRepository(
session_factory=session_factory,
session_factory=sqlite_session_factory,
tenant_id=app.tenant_id,
user=end_user,
app_id=app.id,
triggered_from=WorkflowRunTriggeredFrom.APP_RUN,
),
SQLAlchemyWorkflowNodeExecutionRepository(
session_factory=session_factory,
session_factory=sqlite_session_factory,
tenant_id=app.tenant_id,
user=end_user,
app_id=app.id,
@@ -286,7 +285,11 @@ def test_generate_includes_parent_trace_context_in_extras(
assert repository_tenant_ids == {"workflow": app.tenant_id, "node": app.tenant_id}
def test_resume_delegates_to_generate(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None:
def test_resume_delegates_to_generate(
monkeypatch: pytest.MonkeyPatch,
sqlite_session: Session,
sqlite_session_factory: sessionmaker[Session],
) -> None:
generator = WorkflowAppGenerator()
app, workflow, end_user = _persist_generator_rows(sqlite_session)
mock_generate = MagicMock(return_value="ok")
@@ -300,9 +303,11 @@ def test_resume_delegates_to_generate(monkeypatch: pytest.MonkeyPatch, sqlite_se
stream=False,
)
runtime_state = _runtime_state()
workflow_execution_repository, workflow_node_execution_repository = _repositories(sqlite_session, app, end_user)
workflow_execution_repository, workflow_node_execution_repository = _repositories(
sqlite_session_factory, app, end_user
)
pause_config = PauseStateLayerConfig(
session_factory=sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False),
session_factory=sqlite_session_factory,
state_owner_user_id="owner",
)
@@ -328,7 +333,9 @@ def test_resume_delegates_to_generate(monkeypatch: pytest.MonkeyPatch, sqlite_se
def test_generate_appends_pause_layer_and_forwards_state(
monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
monkeypatch: pytest.MonkeyPatch,
sqlite_session: Session,
sqlite_session_factory: sessionmaker[Session],
) -> None:
generator = WorkflowAppGenerator()
app, workflow, end_user = _persist_generator_rows(sqlite_session)
@@ -360,9 +367,8 @@ def test_generate_appends_pause_layer_and_forwards_state(
staticmethod(get_recording_draft_var_saver_factory),
)
engine = sqlite_session.get_bind()
scoped_session = Session(engine, expire_on_commit=False)
monkeypatch.setattr(app_generator_module.db, "session", scoped_session)
db_session = sqlite_session_factory()
monkeypatch.setattr(app_generator_module.db, "session", db_session)
worker_kwargs: dict[str, object] = {}
@@ -385,7 +391,9 @@ def test_generate_appends_pause_layer_and_forwards_state(
application_generate_entity = _generate_entity(app, workflow, end_user)
graph_runtime_state = _runtime_state()
workflow_execution_repository, workflow_node_execution_repository = _repositories(sqlite_session, app, end_user)
workflow_execution_repository, workflow_node_execution_repository = _repositories(
sqlite_session_factory, app, end_user
)
flask_app = Flask(__name__)
with flask_app.app_context():
@@ -401,7 +409,7 @@ def test_generate_appends_pause_layer_and_forwards_state(
graph_engine_layers=("base-layer",),
graph_runtime_state=graph_runtime_state,
pause_state_config=PauseStateLayerConfig(
session_factory=sessionmaker(bind=engine, expire_on_commit=False),
session_factory=sqlite_session_factory,
state_owner_user_id="owner",
),
)
@@ -416,7 +424,11 @@ def test_generate_appends_pause_layer_and_forwards_state(
assert draft_factory_tenant_ids == [app.tenant_id]
def test_resume_path_runs_worker_with_runtime_state(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None:
def test_resume_path_runs_worker_with_runtime_state(
monkeypatch: pytest.MonkeyPatch,
sqlite_session: Session,
sqlite_session_factory: sessionmaker[Session],
) -> None:
generator = WorkflowAppGenerator()
app, workflow, end_user = _persist_generator_rows(sqlite_session)
workflow_run = WorkflowRun(
@@ -449,13 +461,8 @@ def test_resume_path_runs_worker_with_runtime_state(monkeypatch: pytest.MonkeyPa
MagicMock(side_effect=lambda response, invoke_from: response),
)
engine = sqlite_session.get_bind()
monkeypatch.setattr(app_generator_module.db, "session", Session(engine, expire_on_commit=False))
monkeypatch.setattr(
app_generator_module.session_factory,
"create_session",
lambda: Session(engine, expire_on_commit=False),
)
db_session = sqlite_session_factory()
monkeypatch.setattr(app_generator_module.db, "session", db_session)
runner_instance = MagicMock()
@@ -490,12 +497,14 @@ def test_resume_path_runs_worker_with_runtime_state(monkeypatch: pytest.MonkeyPa
monkeypatch.setattr("core.app.apps.workflow.app_generator.threading.Thread", ImmediateThread)
pause_config = PauseStateLayerConfig(
session_factory=sessionmaker(bind=engine, expire_on_commit=False),
session_factory=sqlite_session_factory,
state_owner_user_id="owner",
)
application_generate_entity = _generate_entity(app, workflow, end_user)
workflow_execution_repository, workflow_node_execution_repository = _repositories(sqlite_session, app, end_user)
workflow_execution_repository, workflow_node_execution_repository = _repositories(
sqlite_session_factory, app, end_user
)
result = generator.resume(
app_model=app,
@@ -2,13 +2,12 @@ from __future__ import annotations
import contextlib
import json
from collections.abc import Iterator
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
from sqlalchemy import inspect
from sqlalchemy.orm import Session, scoped_session, sessionmaker
from sqlalchemy.orm import Session, sessionmaker
from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig
from core.app.apps.exc import GenerateTaskStoppedError
@@ -30,27 +29,22 @@ CREATOR_ID = "00000000-0000-0000-0000-000000000006"
@pytest.fixture
def sqlite_generator_scoped_session(
def sqlite_generator_session(
monkeypatch: pytest.MonkeyPatch,
sqlite_engine,
sqlite_session: Session,
sqlite_session_factory: sessionmaker[Session],
) -> Iterator[scoped_session[Session]]:
"""Adapt the shared SQLite session factory to Flask-SQLAlchemy's scoped session API."""
engine = sqlite_session.get_bind()
request_sessions = scoped_session(sqlite_session_factory)
) -> Session:
"""Expose the shared SQLite engine and real Session through the Flask extension boundary."""
monkeypatch.setattr(
app_generator_module,
"db",
SimpleNamespace(engine=engine, session=request_sessions),
SimpleNamespace(engine=sqlite_engine, session=sqlite_session),
)
try:
yield request_sessions
finally:
request_sessions.remove()
return sqlite_session
def _persist_app(session: Session) -> App:
app = App(
def _app() -> App:
return App(
id=APP_ID,
tenant_id=TENANT_ID,
name="Workflow app",
@@ -66,13 +60,16 @@ def _persist_app(session: Session) -> App:
max_active_requests=None,
created_by=CREATOR_ID,
)
def _persist_app(session: Session) -> App:
app = _app()
session.add(app)
session.commit()
return app
def _persist_workflow(
session: Session,
def _workflow(
*,
workflow_id: str = WORKFLOW_ID,
app_id: str = APP_ID,
@@ -93,13 +90,30 @@ def _persist_workflow(
kind=kind.value,
)
workflow.id = workflow_id
return workflow
def _persist_workflow(
session: Session,
*,
workflow_id: str = WORKFLOW_ID,
app_id: str = APP_ID,
tenant_id: str = TENANT_ID,
kind: WorkflowKind = WorkflowKind.STANDARD,
) -> Workflow:
workflow = _workflow(
workflow_id=workflow_id,
app_id=app_id,
tenant_id=tenant_id,
kind=kind,
)
session.add(workflow)
session.commit()
return workflow
def _persist_end_user(session: Session) -> EndUser:
end_user = EndUser(
def _end_user() -> EndUser:
return EndUser(
id=END_USER_ID,
tenant_id=TENANT_ID,
app_id=APP_ID,
@@ -107,6 +121,10 @@ def _persist_end_user(session: Session) -> EndUser:
name="End user",
session_id="session-id",
)
def _persist_end_user(session: Session) -> EndUser:
end_user = _end_user()
session.add(end_user)
session.commit()
return end_user
@@ -131,9 +149,14 @@ def _persist_snippet(
class TestWorkflowAppGeneratorValidation:
@pytest.mark.usefixtures("sqlite_generator_scoped_session")
def test_generate_stream_joins_worker_after_response_exhaustion(self, monkeypatch: pytest.MonkeyPatch):
@pytest.mark.usefixtures("sqlite_generator_session")
def test_generate_stream_joins_worker_after_response_exhaustion(
self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
):
generator = WorkflowAppGenerator()
app = _persist_app(sqlite_session)
workflow = _persist_workflow(sqlite_session)
user = _persist_end_user(sqlite_session)
worker_thread = Mock()
worker_thread.is_alive.return_value = False
app_config = WorkflowUIBasedAppConfig(
@@ -177,9 +200,9 @@ class TestWorkflowAppGeneratorValidation:
)
managed_stream = generator._generate(
app_model=SimpleNamespace(mode=AppMode.WORKFLOW, tenant_id="tenant"),
workflow=SimpleNamespace(id="workflow-id"),
user=SimpleNamespace(id="user"),
app_model=app,
workflow=workflow,
user=user,
application_generate_entity=application_generate_entity,
invoke_from=InvokeFrom.WEB_APP,
workflow_execution_repository=SimpleNamespace(),
@@ -196,7 +219,7 @@ class TestWorkflowAppGeneratorValidation:
self,
unbound_session: Session,
):
workflow = SimpleNamespace(kind_or_standard="workflow")
workflow = _workflow()
result = WorkflowAppGenerator._ensure_snippet_start_node_in_worker(
session=unbound_session,
@@ -226,7 +249,7 @@ class TestWorkflowAppGeneratorValidation:
):
workflow = _persist_workflow(sqlite_session, kind=WorkflowKind.SNIPPET)
snippet = _persist_snippet(sqlite_session, snippet_id=APP_ID)
injected_workflow = SimpleNamespace(id="workflow-injected")
injected_workflow = _workflow(workflow_id="workflow-injected")
ensure_start_node = Mock(return_value=injected_workflow)
monkeypatch.setattr(
"services.snippet_generate_service.SnippetGenerateService.ensure_start_node_for_worker",
@@ -252,10 +275,10 @@ class TestWorkflowAppGeneratorValidation:
with pytest.raises(ValueError, match="node_id is required"):
generator.single_iteration_generate(
app_model=SimpleNamespace(),
workflow=SimpleNamespace(),
app_model=_app(),
workflow=_workflow(),
node_id="",
user=SimpleNamespace(),
user=_end_user(),
args={"inputs": {}},
streaming=False,
session=sqlite_session,
@@ -263,10 +286,10 @@ class TestWorkflowAppGeneratorValidation:
with pytest.raises(ValueError, match="inputs is required"):
generator.single_iteration_generate(
app_model=SimpleNamespace(),
workflow=SimpleNamespace(),
app_model=_app(),
workflow=_workflow(),
node_id="node",
user=SimpleNamespace(),
user=_end_user(),
args={},
streaming=False,
session=sqlite_session,
@@ -277,16 +300,16 @@ class TestWorkflowAppGeneratorValidation:
with pytest.raises(ValueError, match="node_id is required"):
generator.single_loop_generate(
app_model=SimpleNamespace(),
workflow=SimpleNamespace(),
app_model=_app(),
workflow=_workflow(),
node_id="",
user=SimpleNamespace(),
user=_end_user(),
args=SimpleNamespace(inputs={}),
streaming=False,
session=sqlite_session,
)
@pytest.mark.usefixtures("sqlite_generator_scoped_session")
@pytest.mark.usefixtures("sqlite_generator_session")
def test_single_iteration_generate_includes_trace_session_id_in_extras(
self,
monkeypatch: pytest.MonkeyPatch,
@@ -346,7 +369,7 @@ class TestWorkflowAppGeneratorValidation:
assert len(draft_sessions) == 1
assert draft_sessions[0] is sqlite_session
@pytest.mark.usefixtures("sqlite_generator_scoped_session")
@pytest.mark.usefixtures("sqlite_generator_session")
def test_single_loop_generate_includes_trace_session_id_in_extras(
self,
monkeypatch: pytest.MonkeyPatch,
@@ -408,10 +431,10 @@ class TestWorkflowAppGeneratorValidation:
with pytest.raises(ValueError, match="inputs is required"):
generator.single_loop_generate(
app_model=SimpleNamespace(),
workflow=SimpleNamespace(),
app_model=_app(),
workflow=_workflow(),
node_id="node",
user=SimpleNamespace(),
user=_end_user(),
args=SimpleNamespace(inputs=None),
streaming=False,
session=sqlite_session,
@@ -459,16 +482,16 @@ class TestWorkflowAppGeneratorHandleResponse:
with pytest.raises(GenerateTaskStoppedError):
generator._handle_response(
application_generate_entity=application_generate_entity,
workflow=SimpleNamespace(),
workflow=_workflow(),
queue_manager=SimpleNamespace(),
user=SimpleNamespace(),
user=_end_user(),
draft_var_saver_factory=lambda **kwargs: None,
stream=False,
)
class TestWorkflowAppGeneratorGenerate:
@pytest.mark.usefixtures("sqlite_generator_scoped_session")
@pytest.mark.usefixtures("sqlite_generator_session")
def test_generate_skips_prepare_inputs_when_flag_set(
self,
monkeypatch: pytest.MonkeyPatch,
@@ -590,9 +613,9 @@ class TestWorkflowAppGeneratorResume:
monkeypatch.setattr(generator, "_generate", _fake_generate)
result = generator.resume(
app_model=SimpleNamespace(id="app-id"),
workflow=SimpleNamespace(),
user=SimpleNamespace(id="end-user-id", session_id="session-id"),
app_model=_app(),
workflow=_workflow(),
user=_end_user(),
application_generate_entity=application_generate_entity,
graph_runtime_state=SimpleNamespace(),
workflow_execution_repository=SimpleNamespace(),
@@ -603,7 +626,7 @@ class TestWorkflowAppGeneratorResume:
assert captured_entity is not None
trace_manager = captured_entity.trace_manager
assert isinstance(trace_manager, DummyTraceQueueManager)
assert trace_manager.app_id == "app-id"
assert trace_manager.app_id == APP_ID
assert trace_manager.user_id == "session-id"
def test_resume_preserves_existing_trace_manager(self, monkeypatch: pytest.MonkeyPatch):
@@ -640,9 +663,9 @@ class TestWorkflowAppGeneratorResume:
monkeypatch.setattr(generator, "_generate", _fake_generate)
result = generator.resume(
app_model=SimpleNamespace(id="app-id"),
workflow=SimpleNamespace(),
user=SimpleNamespace(id="end-user-id", session_id="session-id"),
app_model=_app(),
workflow=_workflow(),
user=_end_user(),
application_generate_entity=application_generate_entity,
graph_runtime_state=SimpleNamespace(),
workflow_execution_repository=SimpleNamespace(),
@@ -659,7 +682,6 @@ class TestWorkflowAppGeneratorWorker:
self,
monkeypatch: pytest.MonkeyPatch,
sqlite_session: Session,
sqlite_session_factory: sessionmaker[Session],
):
generator = WorkflowAppGenerator()
_persist_app(sqlite_session)
@@ -679,10 +701,6 @@ class TestWorkflowAppGeneratorWorker:
"core.app.apps.workflow.app_generator.preserve_flask_contexts",
lambda flask_app, context_vars: contextlib.nullcontext(),
)
monkeypatch.setattr(
"core.app.apps.workflow.app_generator.session_factory.create_session",
sqlite_session_factory,
)
monkeypatch.setattr("core.app.apps.workflow.app_generator.WorkflowAppRunner", _Runner)
restore_workflow_run_graph = Mock()
monkeypatch.setattr(generator, "_restore_workflow_run_graph", restore_workflow_run_graph)