diff --git a/api/tests/unit_tests/controllers/service_api/app/test_workflow.py b/api/tests/unit_tests/controllers/service_api/app/test_workflow.py index 7975a935f93..404ac2ac611 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_workflow.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_workflow.py @@ -582,7 +582,7 @@ class TestWorkflowRunApi: handler(api, session=sqlite_session, app_model=app_model, end_user=end_user) def test_sandbox_billing_does_not_gate_default_workflow_run( - self, app: Flask, monkeypatch: pytest.MonkeyPatch + self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session ) -> None: workflow_module = sys.modules["controllers.service_api.app.workflow"] monkeypatch.setattr(workflow_module.dify_config, "BILLING_ENABLED", True) @@ -598,7 +598,7 @@ class TestWorkflowRunApi: with app.test_request_context("/workflows/run", method="POST", json={"inputs": {}}): response = handler( api, - session=Mock(), + session=sqlite_session, app_model=_make_app_model(), end_user=_make_end_user(), ) @@ -609,7 +609,9 @@ class TestWorkflowRunApi: class TestWorkflowRunByIdApi: - def test_rejects_sandbox_plan_with_upgrade_error(self, app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + def test_rejects_sandbox_plan_with_upgrade_error( + self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session + ) -> None: workflow_module = sys.modules["controllers.service_api.app.workflow"] monkeypatch.setattr(workflow_module.dify_config, "BILLING_ENABLED", True) @@ -626,7 +628,7 @@ class TestWorkflowRunByIdApi: with pytest.raises(WorkflowVersionExecutionNotAllowedError) as exc_info: handler( api, - session=Mock(), + session=sqlite_session, app_model=app_model, end_user=_make_end_user(), workflow_id="w1", @@ -660,6 +662,7 @@ class TestWorkflowRunByIdApi: billing_config_enabled: bool, billing_enabled: bool, plan: CloudPlan, + sqlite_session: Session, ) -> None: workflow_module = sys.modules["controllers.service_api.app.workflow"] monkeypatch.setattr(workflow_module.dify_config, "BILLING_ENABLED", billing_config_enabled) @@ -676,7 +679,7 @@ class TestWorkflowRunByIdApi: with app.test_request_context("/workflows/w1/run", method="POST", json={"inputs": {}}): response = handler( api, - session=Mock(), + session=sqlite_session, app_model=app_model, end_user=_make_end_user(), workflow_id="w1", diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_input_guards.py b/api/tests/unit_tests/core/app/apps/agent_app/test_input_guards.py index ec2706a2068..dc14fb9895c 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_input_guards.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_input_guards.py @@ -10,9 +10,9 @@ from __future__ import annotations from types import SimpleNamespace from typing import Any -from unittest.mock import MagicMock import pytest +from sqlalchemy.orm import Session import core.app.features.annotation_reply.annotation_reply as annotation_mod import core.moderation.input_moderation as input_moderation_mod @@ -75,13 +75,13 @@ def _saved_user_query(events: list[Any]) -> str: class TestRunInputGuards: - def test_no_guards_passes_through(self, monkeypatch: pytest.MonkeyPatch): + def test_no_guards_passes_through(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): _patch_moderation(monkeypatch, returns=(False, {}, "hello")) _patch_annotation(monkeypatch, reply=None) qm = _FakeQueueManager() handled, query, annotation_reply = AgentAppGenerator()._run_input_guards( - session=MagicMock(), + session=sqlite_session, application_generate_entity=_make_entity("hello"), app_model=SimpleNamespace(id="app-1"), message=SimpleNamespace(id="msg-1"), @@ -93,13 +93,13 @@ class TestRunInputGuards: assert annotation_reply is None assert qm.events == [] - def test_moderation_override_sanitizes_query(self, monkeypatch: pytest.MonkeyPatch): + def test_moderation_override_sanitizes_query(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): _patch_moderation(monkeypatch, returns=(True, {}, "[redacted]")) _patch_annotation(monkeypatch, reply=None) qm = _FakeQueueManager() handled, query, annotation_reply = AgentAppGenerator()._run_input_guards( - session=MagicMock(), + session=sqlite_session, application_generate_entity=_make_entity("leak my secret"), app_model=SimpleNamespace(id="app-1"), message=SimpleNamespace(id="msg-1"), @@ -111,13 +111,13 @@ class TestRunInputGuards: assert annotation_reply is None assert qm.events == [] - def test_moderation_block_short_circuits(self, monkeypatch: pytest.MonkeyPatch): + def test_moderation_block_short_circuits(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): _patch_moderation(monkeypatch, raises=ModerationError("blocked preset answer")) _patch_annotation(monkeypatch, reply=None) qm = _FakeQueueManager() handled, _, annotation_reply = AgentAppGenerator()._run_input_guards( - session=MagicMock(), + session=sqlite_session, application_generate_entity=_make_entity("forbidden"), app_model=SimpleNamespace(id="app-1"), message=SimpleNamespace(id="msg-1"), @@ -130,13 +130,13 @@ class TestRunInputGuards: assert _answer_text(qm.events) == "blocked preset answer" assert _saved_user_query(qm.events) == "forbidden" - def test_annotation_hit_short_circuits(self, monkeypatch: pytest.MonkeyPatch): + def test_annotation_hit_short_circuits(self, monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): _patch_moderation(monkeypatch, returns=(False, {}, "what is your name")) _patch_annotation(monkeypatch, reply=SimpleNamespace(id="anno-1", content="I am the annotated Iris.")) qm = _FakeQueueManager() handled, _, annotation_reply = AgentAppGenerator()._run_input_guards( - session=MagicMock(), + session=sqlite_session, application_generate_entity=_make_entity("what is your name"), app_model=SimpleNamespace(id="app-1"), message=SimpleNamespace(id="msg-1"), diff --git a/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_generator.py b/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_generator.py index 0dadae0064b..29bfe2b4bb3 100644 --- a/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_generator.py @@ -5,6 +5,7 @@ import logging import pytest from pydantic import ValidationError from pytest_mock import MockerFixture +from sqlalchemy.orm import Session, sessionmaker from core.app.apps.agent_chat.app_generator import AgentChatAppGenerator from core.app.apps.exc import GenerateTaskStoppedError @@ -30,12 +31,12 @@ def generator(mocker: MockerFixture): class TestAgentChatAppGeneratorGenerate: - def test_generate_rejects_blocking_mode(self, generator, mocker: MockerFixture): + def test_generate_rejects_blocking_mode(self, generator, mocker: MockerFixture, sqlite_session: Session): app_model = mocker.MagicMock() user = DummyAccount("user") with pytest.raises(ValueError): generator.generate( - session=mocker.MagicMock(), + session=sqlite_session, app_model=app_model, user=user, args={}, @@ -43,44 +44,45 @@ class TestAgentChatAppGeneratorGenerate: streaming=False, ) - def test_generate_requires_query(self, generator, mocker: MockerFixture): + def test_generate_requires_query(self, generator, mocker: MockerFixture, sqlite_session: Session): app_model = mocker.MagicMock() user = DummyAccount("user") with pytest.raises(ValueError): generator.generate( - session=mocker.MagicMock(), + session=sqlite_session, app_model=app_model, user=user, args={"inputs": {}}, invoke_from=mocker.MagicMock(), ) - def test_generate_rejects_non_string_query(self, generator, mocker: MockerFixture): + def test_generate_rejects_non_string_query(self, generator, mocker: MockerFixture, sqlite_session: Session): app_model = mocker.MagicMock() user = DummyAccount("user") with pytest.raises(ValueError): generator.generate( - session=mocker.MagicMock(), + session=sqlite_session, app_model=app_model, user=user, args={"query": 123, "inputs": {}}, invoke_from=mocker.MagicMock(), ) - def test_generate_override_requires_debugger(self, generator, mocker: MockerFixture): + 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()) with pytest.raises(ValueError): generator.generate( - session=mocker.MagicMock(), + session=sqlite_session, app_model=app_model, user=user, args={"query": "hi", "inputs": {}, "model_config": {"model": {"provider": "p"}}}, invoke_from=InvokeFrom.WEB_APP, ) - def test_generate_success_with_debugger_override(self, generator, mocker: MockerFixture): + 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"}} @@ -154,7 +156,7 @@ class TestAgentChatAppGeneratorGenerate: "files": [{"id": "f1"}], "trace_session_id": "session-1", } - session = mocker.MagicMock() + session = sqlite_session result = generator.generate( session=session, @@ -173,7 +175,7 @@ class TestAgentChatAppGeneratorGenerate: inspect.signature(worker_call.kwargs["target"]).bind(**worker_call.kwargs["kwargs"]) thread_obj.start.assert_called_once() - def test_generate_without_file_config(self, generator, mocker: MockerFixture): + 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"}} @@ -235,7 +237,7 @@ class TestAgentChatAppGeneratorGenerate: ) args = {"query": "hello", "inputs": {"name": "world"}} - session = mocker.MagicMock() + session = sqlite_session result = generator.generate( session=session, @@ -261,7 +263,9 @@ 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): + def test_generate_worker_handles_generate_task_stopped( + self, generator, mocker: MockerFixture, 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()) @@ -269,10 +273,10 @@ class TestAgentChatAppGeneratorWorker: runner = mocker.MagicMock() runner.run.side_effect = GenerateTaskStoppedError() mocker.patch("core.app.apps.agent_chat.app_generator.AgentChatAppRunner", return_value=runner) - session_cm = mocker.MagicMock() - session_cm.__enter__.return_value = mocker.MagicMock() - create_session = mocker.patch("core.app.apps.agent_chat.app_generator.session_factory.create_session") - create_session.return_value = session_cm + 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(), @@ -294,7 +298,9 @@ class TestAgentChatAppGeneratorWorker: Exception("bad"), ], ) - def test_generate_worker_publishes_errors(self, generator, mocker: MockerFixture, error): + def test_generate_worker_publishes_errors( + self, generator, mocker: MockerFixture, error, 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()) @@ -302,10 +308,10 @@ class TestAgentChatAppGeneratorWorker: runner = mocker.MagicMock() runner.run.side_effect = error mocker.patch("core.app.apps.agent_chat.app_generator.AgentChatAppRunner", return_value=runner) - session_cm = mocker.MagicMock() - session_cm.__enter__.return_value = mocker.MagicMock() - create_session = mocker.patch("core.app.apps.agent_chat.app_generator.session_factory.create_session") - create_session.return_value = session_cm + 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(), @@ -319,7 +325,11 @@ class TestAgentChatAppGeneratorWorker: assert queue_manager.publish_error.called def test_generate_worker_logs_value_error_when_debug( - self, generator, mocker: MockerFixture, caplog: pytest.LogCaptureFixture + self, + generator, + mocker: MockerFixture, + caplog: pytest.LogCaptureFixture, + sqlite_session_factory: sessionmaker[Session], ): queue_manager = mocker.MagicMock() generator._get_conversation = mocker.MagicMock(return_value=mocker.MagicMock()) @@ -328,10 +338,10 @@ class TestAgentChatAppGeneratorWorker: runner = mocker.MagicMock() runner.run.side_effect = ValueError("bad") mocker.patch("core.app.apps.agent_chat.app_generator.AgentChatAppRunner", return_value=runner) - session_cm = mocker.MagicMock() - session_cm.__enter__.return_value = mocker.MagicMock() - create_session = mocker.patch("core.app.apps.agent_chat.app_generator.session_factory.create_session") - create_session.return_value = session_cm + mocker.patch( + "core.app.apps.agent_chat.app_generator.session_factory.create_session", + side_effect=sqlite_session_factory, + ) mocker.patch("core.app.apps.agent_chat.app_generator.dify_config", new=mocker.MagicMock(DEBUG=True)) diff --git a/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_node_execution_repository.py b/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_node_execution_repository.py index 8ddba55f6ce..b3ba655eaa4 100644 --- a/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_node_execution_repository.py +++ b/api/tests/unit_tests/extensions/logstore/repositories/test_logstore_workflow_node_execution_repository.py @@ -3,6 +3,7 @@ from typing import cast from unittest.mock import MagicMock, patch import pytest +from sqlalchemy.orm import Session, sessionmaker from extensions.logstore.repositories.logstore_workflow_node_execution_repository import ( LogstoreWorkflowNodeExecutionRepository, @@ -11,7 +12,9 @@ from models.account import Account from models.workflow import WorkflowNodeExecutionTriggeredFrom -def test_save_synchronously_writes_sql_when_dual_write_is_disabled(monkeypatch: pytest.MonkeyPatch) -> None: +def test_save_synchronously_writes_sql_when_dual_write_is_disabled( + monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session] +) -> None: monkeypatch.delenv("LOGSTORE_DUAL_WRITE_ENABLED", raising=False) with ( patch("extensions.logstore.repositories.logstore_workflow_node_execution_repository.AliyunLogStore"), @@ -21,7 +24,7 @@ def test_save_synchronously_writes_sql_when_dual_write_is_disabled(monkeypatch: ) as sql_repository_type, ): repository = LogstoreWorkflowNodeExecutionRepository( - session_factory=MagicMock(), + session_factory=sqlite_session_factory, tenant_id="tenant-1", user=cast(Account, SimpleNamespace(id="account-1")), app_id="app-1", diff --git a/api/tests/unit_tests/services/plugin/test_plugin_migration.py b/api/tests/unit_tests/services/plugin/test_plugin_migration.py index 27b9749bf11..58f7d8cef16 100644 --- a/api/tests/unit_tests/services/plugin/test_plugin_migration.py +++ b/api/tests/unit_tests/services/plugin/test_plugin_migration.py @@ -1,9 +1,12 @@ import json +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest from pytest_mock import MockerFixture +from sqlalchemy.orm import Session +from models.model import App, AppMode from services.plugin.plugin_migration import PluginMigration MIGRATION_MODULE = "services.plugin.plugin_migration" @@ -33,24 +36,28 @@ def test_fetch_latest_package_identifier_calls_marketplace_when_enabled(mocker: assert result == "langgenius/openai:1.0.0@abc" -def test_extract_app_tables_checks_agent_mode_with_its_session(mocker: MockerFixture) -> None: - app = mocker.MagicMock(app_model_config_id=None, mode="chat") - app.is_agent_with_session.return_value = False - apps_result = mocker.MagicMock() - apps_result.all.return_value = [app] - configs_result = mocker.MagicMock() - configs_result.all.return_value = [] - session = mocker.MagicMock() - session.scalars.side_effect = [apps_result, configs_result] - session_context = mocker.MagicMock() - session_context.__enter__.return_value = session - mocker.patch(f"{MIGRATION_MODULE}.Session", return_value=session_context) - mocker.patch(f"{MIGRATION_MODULE}.db") +def test_extract_app_tables_checks_agent_mode_with_its_session(mocker: MockerFixture, sqlite_session: Session) -> None: + app = App( + id="app-1", + tenant_id="tenant-1", + name="Chat app", + description="", + mode=AppMode.CHAT, + icon_type=None, + icon="", + icon_background=None, + enable_site=False, + enable_api=False, + created_by="account-1", + max_active_requests=0, + ) + sqlite_session.add(app) + sqlite_session.commit() + mocker.patch(f"{MIGRATION_MODULE}.db", SimpleNamespace(engine=sqlite_session.get_bind())) result = PluginMigration.extract_app_tables("tenant-1") assert result == [] - app.is_agent_with_session.assert_called_once_with(session=session) class TestHandlePluginInstanceInstall: diff --git a/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py b/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py index 1ad4e070edd..ac9303b2bc4 100644 --- a/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py +++ b/api/tests/unit_tests/services/rag_pipeline/test_pipeline_generate_service.py @@ -3,14 +3,51 @@ from typing import cast import pytest from pytest_mock import MockerFixture +from sqlalchemy.orm import Session from core.app.entities.app_invoke_entities import InvokeFrom -from models.dataset import Dataset, Pipeline +from core.rag.index_processor.constant.index_type import IndexStructureType +from models.dataset import Dataset, Document, Pipeline +from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus from models.model import Account, App, EndUser from services.dataset_ref_service import DatasetRefService from services.rag_pipeline.pipeline_generate_service import PipelineGenerateService +def _make_pipeline(*, tenant_id: str = "tenant-1") -> Pipeline: + pipeline = Pipeline(tenant_id=tenant_id, name="Pipeline", description="") + pipeline.id = "pipeline-1" + return pipeline + + +def _make_dataset(*, dataset_id: str = "dataset-1", tenant_id: str = "tenant-1") -> Dataset: + return Dataset( + id=dataset_id, + tenant_id=tenant_id, + name="Dataset", + created_by="user-1", + pipeline_id="pipeline-1", + ) + + +def _make_document( + *, document_id: str = "doc-1", dataset_id: str = "dataset-1", tenant_id: str = "tenant-1" +) -> Document: + return Document( + id=document_id, + tenant_id=tenant_id, + dataset_id=dataset_id, + position=1, + data_source_type=DataSourceType.LOCAL_FILE, + batch="batch", + name="Document", + created_from=DocumentCreatedFrom.API, + created_by="user-1", + indexing_status=IndexingStatus.COMPLETED, + doc_form=IndexStructureType.PARAGRAPH_INDEX, + ) + + def test_get_max_active_requests_uses_smallest_non_zero_limit(mocker: MockerFixture) -> None: mocker.patch("services.rag_pipeline.pipeline_generate_service.dify_config.APP_DEFAULT_ACTIVE_REQUESTS", 5) mocker.patch("services.rag_pipeline.pipeline_generate_service.dify_config.APP_MAX_ACTIVE_REQUESTS", 3) @@ -41,14 +78,14 @@ def test_get_max_active_requests_returns_zero_when_all_unlimited(mocker: MockerF (InvokeFrom.DEBUGGER, SimpleNamespace(id="wf-1"), None), ], ) -def test_get_workflow(mocker: MockerFixture, invoke_from, workflow, expected_error) -> None: +def test_get_workflow(mocker: MockerFixture, invoke_from, workflow, expected_error, sqlite_session: Session) -> None: rag_pipeline_service_cls = mocker.patch("services.rag_pipeline.pipeline_generate_service.RagPipelineService") rag_pipeline_service = rag_pipeline_service_cls.return_value rag_pipeline_service.get_draft_workflow.return_value = workflow rag_pipeline_service.get_published_workflow.return_value = workflow pipeline = cast(Pipeline, SimpleNamespace(id="pipeline-1")) - session = mocker.Mock() + session = sqlite_session if expected_error: with pytest.raises(ValueError, match=expected_error): @@ -58,7 +95,9 @@ def test_get_workflow(mocker: MockerFixture, invoke_from, workflow, expected_err assert result == workflow -def test_generate_updates_document_status_and_returns_event_stream(mocker: MockerFixture) -> None: +def test_generate_updates_document_status_and_returns_event_stream( + mocker: MockerFixture, sqlite_session: Session +) -> None: dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-1")) pipeline = cast( Pipeline, @@ -70,7 +109,6 @@ def test_generate_updates_document_status_and_returns_event_stream(mocker: Mocke ) user = cast(Account | EndUser, SimpleNamespace(id="user-1")) args = {"original_document_id": "doc-1", "query": "hello"} - session_mock = mocker.Mock() mocker.patch.object(PipelineGenerateService, "_get_workflow", return_value=SimpleNamespace(id="wf-1")) update_status_mock = mocker.patch.object(PipelineGenerateService, "update_document_status") @@ -86,7 +124,7 @@ def test_generate_updates_document_status_and_returns_event_stream(mocker: Mocke args=args, invoke_from=InvokeFrom.WEB_APP, streaming=True, - session=session_mock, + session=sqlite_session, ) assert result == "stream-events" @@ -94,11 +132,11 @@ def test_generate_updates_document_status_and_returns_event_stream(mocker: Mocke assert document_ref.dataset.tenant_id == "tenant-1" assert document_ref.dataset.dataset_id == "dataset-1" assert document_ref.document_id == "doc-1" - update_status_mock.assert_called_once_with(document_ref, session=session_mock) - assert generator_instance.generate.call_args.kwargs["session"] is session_mock + update_status_mock.assert_called_once_with(document_ref, session=sqlite_session) + assert generator_instance.generate.call_args.kwargs["session"] is sqlite_session -def test_generate_rejects_pipeline_dataset_from_another_tenant(mocker: MockerFixture) -> None: +def test_generate_rejects_pipeline_dataset_from_another_tenant(mocker: MockerFixture, sqlite_session: Session) -> None: dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-2")) pipeline = cast( Pipeline, @@ -117,7 +155,7 @@ def test_generate_rejects_pipeline_dataset_from_another_tenant(mocker: MockerFix user=cast(Account, SimpleNamespace(id="user-1")), args={"original_document_id": "doc-1"}, invoke_from=InvokeFrom.WEB_APP, - session=mocker.Mock(), + session=sqlite_session, ) update_status_mock.assert_not_called() @@ -125,18 +163,13 @@ def test_generate_rejects_pipeline_dataset_from_another_tenant(mocker: MockerFix def test_generate_rejects_original_document_outside_pipeline_dataset_before_dispatch( mocker: MockerFixture, + sqlite_session: Session, ) -> None: - dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-1")) - pipeline = cast( - Pipeline, - SimpleNamespace( - id="pipeline-1", - tenant_id="tenant-1", - retrieve_dataset=mocker.Mock(return_value=dataset), - ), - ) - session = mocker.Mock() - session.scalar.return_value = None + dataset = _make_dataset() + pipeline = _make_pipeline() + outside_document = _make_document(document_id="foreign-doc", dataset_id="other-dataset", tenant_id="tenant-2") + sqlite_session.add_all([dataset, pipeline, outside_document]) + sqlite_session.commit() mocker.patch.object(PipelineGenerateService, "_get_workflow", return_value=SimpleNamespace(id="wf-1")) generator_cls = mocker.patch("services.rag_pipeline.pipeline_generate_service.PipelineGenerator") @@ -146,29 +179,25 @@ def test_generate_rejects_original_document_outside_pipeline_dataset_before_disp user=cast(Account, SimpleNamespace(id="user-1")), args={"original_document_id": "foreign-doc"}, invoke_from=InvokeFrom.PUBLISHED_PIPELINE, - session=session, + session=sqlite_session, ) - statement = session.scalar.call_args.args[0] - assert {"foreign-doc", "dataset-1", "tenant-1"} <= set(statement.compile().params.values()) + sqlite_session.refresh(outside_document) + assert outside_document.indexing_status == IndexingStatus.COMPLETED generator_cls.assert_not_called() -def test_update_document_status_updates_existing_document(mocker: MockerFixture) -> None: - document = SimpleNamespace(indexing_status="completed") - dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-1")) +def test_update_document_status_updates_existing_document(sqlite_session: Session) -> None: + document = _make_document() + dataset = _make_dataset() + sqlite_session.add_all([dataset, document]) + sqlite_session.commit() dataset_ref = DatasetRefService.create_dataset_ref(dataset) document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, "doc-1") - session_mock = mocker.Mock() - get_document_mock = mocker.patch.object(DatasetRefService, "get_document_by_ref", return_value=document) - add_mock = session_mock.add + PipelineGenerateService.update_document_status(document_ref, session=sqlite_session) - PipelineGenerateService.update_document_status(document_ref, session=session_mock) - - assert document.indexing_status == "waiting" - get_document_mock.assert_called_once_with(document_ref, session=session_mock) - add_mock.assert_called_once_with(document) + assert document.indexing_status == IndexingStatus.WAITING @pytest.mark.parametrize( @@ -179,42 +208,28 @@ def test_update_document_status_updates_existing_document(mocker: MockerFixture) ], ) def test_update_document_status_rejects_document_outside_owner( - mocker: MockerFixture, document_tenant_id: str, document_dataset_id: str, + sqlite_session: Session, ) -> None: dataset = cast(Dataset, SimpleNamespace(id="dataset-1", tenant_id="tenant-1")) dataset_ref = DatasetRefService.create_dataset_ref(dataset) document_ref = DatasetRefService.create_document_ref_from_id(dataset_ref, "doc-1") - outside_document = SimpleNamespace( - id="doc-1", - tenant_id=document_tenant_id, - dataset_id=document_dataset_id, - indexing_status="completed", - ) - session_mock = mocker.Mock() - - def resolve_document(statement): - params = set(statement.compile().params.values()) - outside_owner = {outside_document.id, outside_document.dataset_id, outside_document.tenant_id} - return outside_document if outside_owner <= params else None - - session_mock.scalar.side_effect = resolve_document - add_mock = session_mock.add + outside_document = _make_document(tenant_id=document_tenant_id, dataset_id=document_dataset_id) + sqlite_session.add(outside_document) + sqlite_session.commit() with pytest.raises(ValueError, match="Pipeline document not found"): - PipelineGenerateService.update_document_status(document_ref, session=session_mock) + PipelineGenerateService.update_document_status(document_ref, session=sqlite_session) - statement = session_mock.scalar.call_args.args[0] - assert {"doc-1", "dataset-1", "tenant-1"} <= set(statement.compile().params.values()) - assert outside_document.indexing_status == "completed" - add_mock.assert_not_called() + sqlite_session.refresh(outside_document) + assert outside_document.indexing_status == IndexingStatus.COMPLETED # --- generate_single_iteration --- -def test_generate_single_iteration_delegates(mocker: MockerFixture) -> None: +def test_generate_single_iteration_delegates(mocker: MockerFixture, sqlite_session: Session) -> None: mocker.patch.object(PipelineGenerateService, "_get_workflow", return_value=SimpleNamespace(id="wf-1")) generator_cls = mocker.patch("services.rag_pipeline.pipeline_generate_service.PipelineGenerator") @@ -224,7 +239,7 @@ def test_generate_single_iteration_delegates(mocker: MockerFixture) -> None: pipeline = cast(Pipeline, SimpleNamespace(id="p1")) user = cast(Account, SimpleNamespace(id="u1")) - session = mocker.Mock() + session = sqlite_session result = PipelineGenerateService.generate_single_iteration(pipeline, user, "node-1", {"key": "val"}, session) @@ -236,7 +251,7 @@ def test_generate_single_iteration_delegates(mocker: MockerFixture) -> None: # --- generate_single_loop --- -def test_generate_single_loop_delegates(mocker: MockerFixture) -> None: +def test_generate_single_loop_delegates(mocker: MockerFixture, sqlite_session: Session) -> None: mocker.patch.object(PipelineGenerateService, "_get_workflow", return_value=SimpleNamespace(id="wf-1")) generator_cls = mocker.patch("services.rag_pipeline.pipeline_generate_service.PipelineGenerator") @@ -246,7 +261,7 @@ def test_generate_single_loop_delegates(mocker: MockerFixture) -> None: pipeline = cast(Pipeline, SimpleNamespace(id="p1")) user = cast(Account, SimpleNamespace(id="u1")) - session = mocker.Mock() + session = sqlite_session result = PipelineGenerateService.generate_single_loop(pipeline, user, "node-1", {"key": "val"}, session) diff --git a/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service_additional.py b/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service_additional.py index 924d272f686..36cdac0742e 100644 --- a/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service_additional.py +++ b/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service_additional.py @@ -362,6 +362,7 @@ class TestBuildWorkflowEventStream: def test_build_workflow_event_stream_should_emit_ping_and_terminal_snapshot_event( self, monkeypatch: pytest.MonkeyPatch, + message_session_maker: sessionmaker[Session], ) -> None: workflow_run = _build_workflow_run(status=WorkflowExecutionStatus.PAUSED) topic = _Topic(_StaticSubscription()) @@ -409,7 +410,7 @@ class TestBuildWorkflowEventStream: workflow_run=workflow_run, tenant_id="tenant-1", app_id="app-1", - session_maker=MagicMock(), + session_maker=message_session_maker, ) ) @@ -424,6 +425,7 @@ class TestBuildWorkflowEventStream: def test_build_workflow_event_stream_should_emit_periodic_ping_and_stop_after_idle_timeout( self, monkeypatch: pytest.MonkeyPatch, + message_session_maker: sessionmaker[Session], ) -> None: workflow_run = _build_workflow_run(status=WorkflowExecutionStatus.RUNNING) topic = _Topic(_StaticSubscription()) @@ -463,7 +465,7 @@ class TestBuildWorkflowEventStream: workflow_run=workflow_run, tenant_id="tenant-1", app_id="app-1", - session_maker=MagicMock(), + session_maker=message_session_maker, idle_timeout=20.0, ping_interval=5.0, ) @@ -475,6 +477,7 @@ class TestBuildWorkflowEventStream: def test_build_workflow_event_stream_should_exit_when_buffer_done_and_empty( self, monkeypatch: pytest.MonkeyPatch, + message_session_maker: sessionmaker[Session], ) -> None: workflow_run = _build_workflow_run(status=WorkflowExecutionStatus.RUNNING) topic = _Topic(_StaticSubscription()) @@ -505,7 +508,7 @@ class TestBuildWorkflowEventStream: workflow_run=workflow_run, tenant_id="tenant-1", app_id="app-1", - session_maker=MagicMock(), + session_maker=message_session_maker, ) ) @@ -515,6 +518,7 @@ class TestBuildWorkflowEventStream: def test_build_workflow_event_stream_should_continue_when_pause_loading_fails( self, monkeypatch: pytest.MonkeyPatch, + message_session_maker: sessionmaker[Session], ) -> None: workflow_run = _build_workflow_run(status=WorkflowExecutionStatus.PAUSED) topic = _Topic(_StaticSubscription()) @@ -545,7 +549,7 @@ class TestBuildWorkflowEventStream: workflow_run=workflow_run, tenant_id="tenant-1", app_id="app-1", - session_maker=MagicMock(), + session_maker=message_session_maker, ) )