mirror of
https://github.com/langgenius/dify.git
synced 2026-09-19 10:11:30 +08:00
test: move RAG pipeline workflow coverage to unit tests (#38937)
This commit is contained in:
Asuka Minato
+193
-111
@@ -1,40 +1,51 @@
|
||||
"""RAG pipeline workflow controller serialization tests.
|
||||
|
||||
Handlers that own transactions run against real SQLite sessions so response
|
||||
DTOs must be materialized before those transaction contexts close.
|
||||
"""
|
||||
"""Unit coverage for RAG workflow controllers using real models and disposable SQLite state."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Iterator
|
||||
from datetime import datetime
|
||||
from inspect import unwrap as unwrap_all
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import PropertyMock, patch
|
||||
from unittest.mock import MagicMock, patch
|
||||
from uuid import UUID
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy.engine import Engine
|
||||
from sqlalchemy import Engine
|
||||
from sqlalchemy.orm import Session
|
||||
from werkzeug.exceptions import Forbidden
|
||||
|
||||
from controllers.console.datasets.rag_pipeline import rag_pipeline_workflow as module
|
||||
from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError
|
||||
from models.account import Account, TenantAccountRole
|
||||
from models.dataset import Pipeline
|
||||
from models.engine import db
|
||||
from models.tools import WorkflowToolProvider
|
||||
from models.workflow import Workflow, WorkflowType
|
||||
from services.errors.llm import InvokeRateLimitError
|
||||
from services.rag_pipeline.rag_pipeline import RagPipelineService
|
||||
|
||||
DEFAULT_WORKFLOW_TENANT_ID = "00000000-0000-0000-0000-000000000001"
|
||||
DEFAULT_WORKFLOW_APP_ID = "00000000-0000-0000-0000-000000000002"
|
||||
DEFAULT_WORKFLOW_CREATED_BY = "00000000-0000-0000-0000-000000000003"
|
||||
DEFAULT_WORKFLOW_ID = "00000000-0000-0000-0000-000000000004"
|
||||
|
||||
|
||||
def _make_workflow(**overrides):
|
||||
workflow = SimpleNamespace(
|
||||
id="workflow-1",
|
||||
graph_dict={"nodes": [], "edges": []},
|
||||
features_dict={"file_upload": {"enabled": False}},
|
||||
unique_hash="hash-1",
|
||||
version="1",
|
||||
def _make_workflow(**overrides: object) -> Workflow:
|
||||
workflow = Workflow(
|
||||
id=DEFAULT_WORKFLOW_ID,
|
||||
tenant_id=DEFAULT_WORKFLOW_TENANT_ID,
|
||||
app_id=DEFAULT_WORKFLOW_APP_ID,
|
||||
type=WorkflowType.WORKFLOW,
|
||||
version=Workflow.VERSION_DRAFT,
|
||||
marked_name="Release 1",
|
||||
marked_comment="Initial release",
|
||||
created_by_account=SimpleNamespace(id="user-1", name="Alice", email="alice@example.com"),
|
||||
graph=json.dumps({"nodes": [], "edges": []}),
|
||||
features=json.dumps({"file_upload": {"enabled": False}}),
|
||||
created_by=DEFAULT_WORKFLOW_CREATED_BY,
|
||||
created_at=datetime(2024, 1, 1, 12, 0, 0),
|
||||
updated_by_account=None,
|
||||
updated_by=None,
|
||||
updated_at=datetime(2024, 1, 1, 12, 1, 0),
|
||||
tool_published=False,
|
||||
environment_variables=[],
|
||||
conversation_variables=[],
|
||||
rag_pipeline_variables=[],
|
||||
@@ -46,137 +57,130 @@ def _make_workflow(**overrides):
|
||||
|
||||
def _account() -> Account:
|
||||
account = Account(name="Alice", email="alice@example.com")
|
||||
account.id = "user-1"
|
||||
account.id = DEFAULT_WORKFLOW_CREATED_BY
|
||||
account.role = TenantAccountRole.EDITOR
|
||||
return account
|
||||
|
||||
|
||||
def _pipeline() -> Pipeline:
|
||||
pipeline = Pipeline(tenant_id="tenant-1", name="Pipeline", description="desc")
|
||||
pipeline.id = "pipeline-1"
|
||||
pipeline = Pipeline(tenant_id=DEFAULT_WORKFLOW_TENANT_ID, name="Pipeline", description="desc")
|
||||
pipeline.id = DEFAULT_WORKFLOW_APP_ID
|
||||
return pipeline
|
||||
|
||||
|
||||
def test_draft_rag_pipeline_workflow_get_serializes_response_model(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def _persist_workflow(workflow: Workflow) -> None:
|
||||
db.session.add(workflow)
|
||||
db.session.commit()
|
||||
db.session.expunge(workflow)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def database_app() -> Iterator[Flask]:
|
||||
app = Flask(__name__)
|
||||
app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:"
|
||||
db.init_app(app)
|
||||
|
||||
with app.app_context():
|
||||
Account.__table__.create(db.engine)
|
||||
WorkflowToolProvider.__table__.create(db.engine)
|
||||
Workflow.__table__.create(db.engine)
|
||||
db.session.add(_account())
|
||||
db.session.commit()
|
||||
|
||||
try:
|
||||
yield app
|
||||
finally:
|
||||
db.session.remove()
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("database_app")
|
||||
def test_draft_rag_pipeline_workflow_get_serializes_response_model() -> None:
|
||||
workflow = _make_workflow()
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"RagPipelineService",
|
||||
lambda *_args, **_kwargs: SimpleNamespace(get_draft_workflow=lambda **_kwargs: workflow),
|
||||
)
|
||||
expected_hash = workflow.unique_hash
|
||||
_persist_workflow(workflow)
|
||||
|
||||
api = module.DraftRagPipelineApi()
|
||||
handler = unwrap_all(api.get)
|
||||
|
||||
response = handler(api, _pipeline())
|
||||
|
||||
assert response["id"] == "workflow-1"
|
||||
assert response["id"] == DEFAULT_WORKFLOW_ID
|
||||
assert response["graph"] == {"nodes": [], "edges": []}
|
||||
assert response["features"] == {"file_upload": {"enabled": False}}
|
||||
assert response["hash"] == "hash-1"
|
||||
assert response["created_by"] == {"id": "user-1", "name": "Alice", "email": "alice@example.com"}
|
||||
assert response["hash"] == expected_hash
|
||||
assert response["created_by"] == {
|
||||
"id": DEFAULT_WORKFLOW_CREATED_BY,
|
||||
"name": "Alice",
|
||||
"email": "alice@example.com",
|
||||
}
|
||||
assert response["updated_by"] is None
|
||||
assert response["created_at"] == int(datetime(2024, 1, 1, 12, 0, 0).timestamp())
|
||||
assert response["updated_at"] == int(datetime(2024, 1, 1, 12, 1, 0).timestamp())
|
||||
|
||||
|
||||
def test_published_rag_pipeline_workflows_serialize_items_before_session_closes(
|
||||
app, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine
|
||||
database_app: Flask,
|
||||
) -> None:
|
||||
api = module.PublishedAllRagPipelineApi()
|
||||
handler = unwrap_all(api.get)
|
||||
session_state: dict[str, Session] = {}
|
||||
workflow = _make_workflow(version="1")
|
||||
_persist_workflow(workflow)
|
||||
pipeline = _pipeline()
|
||||
pipeline.workflow_id = DEFAULT_WORKFLOW_ID
|
||||
|
||||
base_workflow = _make_workflow()
|
||||
with database_app.test_request_context(
|
||||
"/rag/pipelines/pipeline-1/workflows",
|
||||
method="GET",
|
||||
query_string={"page": 1, "limit": 10, "user_id": "", "named_only": "false"},
|
||||
):
|
||||
response = handler(api, _account(), pipeline=pipeline)
|
||||
|
||||
class _Workflow:
|
||||
def __getattr__(self, name: str):
|
||||
assert session_state["session"].in_transaction() is True
|
||||
return getattr(base_workflow, name)
|
||||
|
||||
def _get_all_published_workflow(**kwargs):
|
||||
session_state["session"] = kwargs["session"]
|
||||
return [_Workflow()], False
|
||||
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"RagPipelineService",
|
||||
lambda *_args, **_kwargs: SimpleNamespace(get_all_published_workflow=_get_all_published_workflow),
|
||||
)
|
||||
|
||||
with Session(sqlite_engine) as request_session:
|
||||
monkeypatch.setattr(module, "db", SimpleNamespace(engine=sqlite_engine, session=lambda: request_session))
|
||||
with app.test_request_context(
|
||||
"/rag/pipelines/pipeline-1/workflows",
|
||||
method="GET",
|
||||
query_string={"page": 1, "limit": 10, "user_id": "", "named_only": "false"},
|
||||
):
|
||||
response = handler(api, _account(), pipeline=_pipeline())
|
||||
|
||||
assert session_state["session"].in_transaction() is False
|
||||
assert response["items"][0]["id"] == "workflow-1"
|
||||
assert response["items"][0]["id"] == DEFAULT_WORKFLOW_ID
|
||||
assert response["page"] == 1
|
||||
assert response["limit"] == 10
|
||||
assert response["has_more"] is False
|
||||
|
||||
|
||||
def test_rag_pipeline_workflow_patch_serializes_response_model(
|
||||
app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine
|
||||
database_app: Flask,
|
||||
) -> None:
|
||||
workflow = _make_workflow(marked_name="Updated release")
|
||||
captured_session: dict[str, Session] = {}
|
||||
|
||||
def _update_workflow(**kwargs):
|
||||
captured_session["session"] = kwargs["session"]
|
||||
assert kwargs["session"].in_transaction() is True
|
||||
return workflow
|
||||
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"RagPipelineService",
|
||||
lambda *_args, **_kwargs: SimpleNamespace(update_workflow=_update_workflow),
|
||||
)
|
||||
expected_hash = workflow.unique_hash
|
||||
_persist_workflow(workflow)
|
||||
payload: dict[str, object] = {"marked_name": "Updated release"}
|
||||
|
||||
api = module.RagPipelineByIdApi()
|
||||
handler = unwrap_all(api.patch)
|
||||
|
||||
with Session(sqlite_engine) as request_session:
|
||||
monkeypatch.setattr(module, "db", SimpleNamespace(engine=sqlite_engine, session=lambda: request_session))
|
||||
with (
|
||||
app.test_request_context("/rag/pipelines/pipeline-1/workflows/workflow-1", method="PATCH", json=payload),
|
||||
patch.object(type(module.console_ns), "payload", new_callable=PropertyMock, return_value=payload),
|
||||
):
|
||||
response = handler(
|
||||
api,
|
||||
_account(),
|
||||
pipeline=_pipeline(),
|
||||
workflow_id="workflow-1",
|
||||
)
|
||||
with database_app.test_request_context(
|
||||
f"/rag/pipelines/{DEFAULT_WORKFLOW_APP_ID}/workflows/{DEFAULT_WORKFLOW_ID}", method="PATCH", json=payload
|
||||
):
|
||||
response = handler(
|
||||
api,
|
||||
_account(),
|
||||
pipeline=_pipeline(),
|
||||
workflow_id=DEFAULT_WORKFLOW_ID,
|
||||
)
|
||||
|
||||
assert captured_session["session"].in_transaction() is False
|
||||
assert response["id"] == "workflow-1"
|
||||
assert response["id"] == DEFAULT_WORKFLOW_ID
|
||||
assert response["marked_name"] == "Updated release"
|
||||
assert response["hash"] == "hash-1"
|
||||
assert response["hash"] == expected_hash
|
||||
|
||||
|
||||
def test_default_rag_pipeline_block_configs_serializes_root_response(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
@pytest.mark.usefixtures("database_app")
|
||||
def test_default_rag_pipeline_block_configs_serializes_root_response() -> None:
|
||||
block_configs = [{"type": "start", "config": {"title": "Start"}}]
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"RagPipelineService",
|
||||
lambda *_args, **_kwargs: SimpleNamespace(get_default_block_configs=lambda: block_configs),
|
||||
)
|
||||
|
||||
api = module.DefaultRagPipelineBlockConfigsApi()
|
||||
handler = unwrap_all(api.get)
|
||||
|
||||
response = handler(api, _pipeline())
|
||||
with patch.object(RagPipelineService, "get_default_block_configs", return_value=block_configs):
|
||||
response = handler(api, _pipeline())
|
||||
|
||||
assert response == block_configs
|
||||
|
||||
|
||||
def test_draft_rag_pipeline_second_step_parameters_serializes_variables(app, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def test_draft_rag_pipeline_second_step_parameters_serializes_variables(database_app: Flask) -> None:
|
||||
variables = [
|
||||
{
|
||||
"belong_to_node_id": "shared",
|
||||
@@ -187,36 +191,114 @@ def test_draft_rag_pipeline_second_step_parameters_serializes_variables(app, mon
|
||||
"required": True,
|
||||
}
|
||||
]
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"RagPipelineService",
|
||||
lambda *_args, **_kwargs: SimpleNamespace(get_second_step_parameters=lambda **_kwargs: variables),
|
||||
)
|
||||
|
||||
api = module.DraftRagPipelineSecondStepApi()
|
||||
handler = unwrap_all(api.get)
|
||||
|
||||
with app.test_request_context("/?node_id=node-1"):
|
||||
with (
|
||||
database_app.test_request_context("/?node_id=node-1"),
|
||||
patch.object(RagPipelineService, "get_second_step_parameters", return_value=variables),
|
||||
):
|
||||
response = handler(api, _pipeline())
|
||||
|
||||
assert response["variables"] == variables
|
||||
|
||||
|
||||
def test_rag_pipeline_recommended_plugins_serializes_known_envelope(app, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def test_rag_pipeline_recommended_plugins_serializes_known_envelope(database_app: Flask) -> None:
|
||||
recommended_plugins = {
|
||||
"installed_recommended_plugins": [{"name": "Dify Extractor", "meta": {"version": "1.0.0"}}],
|
||||
"uninstalled_recommended_plugins": [{"plugin_id": "langgenius/notion_datasource"}],
|
||||
}
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"RagPipelineService",
|
||||
lambda *_args, **_kwargs: SimpleNamespace(get_recommended_plugins=lambda *_args: recommended_plugins),
|
||||
)
|
||||
|
||||
api = module.RagPipelineRecommendedPluginApi()
|
||||
handler = unwrap_all(api.get)
|
||||
|
||||
with app.test_request_context("/?type=tool"):
|
||||
response = handler(api, "tenant-1", _account())
|
||||
with (
|
||||
database_app.test_request_context("/?type=tool"),
|
||||
patch.object(RagPipelineService, "get_recommended_plugins", return_value=recommended_plugins),
|
||||
):
|
||||
response = handler(api, DEFAULT_WORKFLOW_TENANT_ID, _account())
|
||||
|
||||
assert response == recommended_plugins
|
||||
|
||||
|
||||
def test_rag_pipeline_transform_rejects_read_only_member(app: Flask, sqlite_engine: Engine) -> None:
|
||||
account = _account()
|
||||
account.role = TenantAccountRole.NORMAL
|
||||
api = module.RagPipelineTransformApi()
|
||||
handler = unwrap_all(api.post)
|
||||
|
||||
with (
|
||||
Session(sqlite_engine) as session,
|
||||
app.test_request_context("/"),
|
||||
pytest.raises(Forbidden),
|
||||
):
|
||||
handler(api, session, account, UUID("44444444-4444-4444-4444-444444444444"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("api_type", "payload"),
|
||||
[
|
||||
(
|
||||
module.DraftRagPipelineRunApi,
|
||||
{"inputs": {}, "datasource_type": "x", "datasource_info_list": [], "start_node_id": "node-1"},
|
||||
),
|
||||
(
|
||||
module.PublishedRagPipelineRunApi,
|
||||
{
|
||||
"inputs": {},
|
||||
"datasource_type": "x",
|
||||
"datasource_info_list": [],
|
||||
"start_node_id": "node-1",
|
||||
"response_mode": "blocking",
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_rag_pipeline_run_uses_sqlite_session(
|
||||
app: Flask,
|
||||
sqlite_engine: Engine,
|
||||
api_type: type,
|
||||
payload: dict[str, object],
|
||||
) -> None:
|
||||
api = api_type()
|
||||
handler = unwrap_all(api.post)
|
||||
pipeline = _pipeline()
|
||||
|
||||
with (
|
||||
Session(sqlite_engine) as session,
|
||||
app.test_request_context("/", json=payload),
|
||||
patch.object(module, "load_rag_pipeline", return_value=pipeline) as load_pipeline,
|
||||
patch.object(module.PipelineGenerateService, "generate", return_value=MagicMock()) as generate,
|
||||
patch.object(module.helper, "compact_generate_response", return_value={"ok": True}),
|
||||
):
|
||||
response = handler(api, session, _account(), pipeline.id)
|
||||
|
||||
assert response == {"ok": True}
|
||||
load_pipeline.assert_called_once_with(session, pipeline.id)
|
||||
assert generate.call_args.kwargs["session"] is session
|
||||
assert session.get_bind() is sqlite_engine
|
||||
|
||||
|
||||
@pytest.mark.parametrize("api_type", [module.DraftRagPipelineRunApi, module.PublishedRagPipelineRunApi])
|
||||
def test_rag_pipeline_run_translates_rate_limit(
|
||||
app: Flask,
|
||||
sqlite_engine: Engine,
|
||||
api_type: type,
|
||||
) -> None:
|
||||
payload = {
|
||||
"inputs": {},
|
||||
"datasource_type": "x",
|
||||
"datasource_info_list": [],
|
||||
"start_node_id": "node-1",
|
||||
}
|
||||
api = api_type()
|
||||
handler = unwrap_all(api.post)
|
||||
pipeline = _pipeline()
|
||||
|
||||
with (
|
||||
Session(sqlite_engine) as session,
|
||||
app.test_request_context("/", json=payload),
|
||||
patch.object(module, "load_rag_pipeline", return_value=pipeline),
|
||||
patch.object(module.PipelineGenerateService, "generate", side_effect=InvokeRateLimitError("limit")),
|
||||
pytest.raises(InvokeRateLimitHttpError),
|
||||
):
|
||||
handler(api, session, _account(), pipeline.id)
|
||||
|
||||
+45
-232
@@ -1,8 +1,10 @@
|
||||
"""Testcontainers integration tests for rag_pipeline_workflow controller endpoints."""
|
||||
"""Unit tests for rag_pipeline_workflow controller endpoints."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from inspect import unwrap
|
||||
from typing import TypedDict, Unpack
|
||||
@@ -11,19 +13,19 @@ from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy.orm import Session
|
||||
from sqlalchemy import Engine
|
||||
from sqlalchemy.orm import Session, scoped_session, sessionmaker
|
||||
from werkzeug.exceptions import BadRequest, Forbidden, HTTPException, NotFound
|
||||
|
||||
import models.workflow as workflow_models
|
||||
import services
|
||||
from controllers.console import console_ns
|
||||
from controllers.console.app.error import DraftWorkflowNotExist, DraftWorkflowNotSync
|
||||
from controllers.console.datasets.rag_pipeline import rag_pipeline_workflow as workflow_controller
|
||||
from controllers.console.datasets.rag_pipeline.rag_pipeline_workflow import (
|
||||
DefaultRagPipelineBlockConfigApi,
|
||||
DraftRagPipelineApi,
|
||||
DraftRagPipelineRunApi,
|
||||
PublishedAllRagPipelineApi,
|
||||
PublishedRagPipelineApi,
|
||||
PublishedRagPipelineRunApi,
|
||||
RagPipelineByIdApi,
|
||||
RagPipelineDatasourceVariableApi,
|
||||
RagPipelineDraftNodeRunApi,
|
||||
@@ -32,11 +34,9 @@ from controllers.console.datasets.rag_pipeline.rag_pipeline_workflow import (
|
||||
RagPipelineDraftWorkflowRestoreApi,
|
||||
RagPipelineRecommendedPluginApi,
|
||||
RagPipelineTaskStopApi,
|
||||
RagPipelineTransformApi,
|
||||
RagPipelineWorkflowLastRunApi,
|
||||
RagPipelineWorkflowRunNodeExecutionListApi,
|
||||
)
|
||||
from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError
|
||||
from graphon.enums import WorkflowNodeExecutionStatus
|
||||
from libs.datetime_utils import naive_utc_now
|
||||
from models.account import Account, TenantAccountRole
|
||||
@@ -44,7 +44,6 @@ from models.dataset import Pipeline
|
||||
from models.enums import CreatorUserRole
|
||||
from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom
|
||||
from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError
|
||||
from services.errors.llm import InvokeRateLimitError
|
||||
|
||||
DEFAULT_WORKFLOW_TENANT_ID = "00000000-0000-0000-0000-000000000001"
|
||||
DEFAULT_WORKFLOW_APP_ID = "00000000-0000-0000-0000-000000000002"
|
||||
@@ -52,6 +51,38 @@ DEFAULT_WORKFLOW_CREATED_BY = "00000000-0000-0000-0000-000000000003"
|
||||
type WorkflowVariablePayload = dict[str, object]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SQLiteDatabase:
|
||||
"""Expose the concrete SQLite engine and scoped session interface used by controller code."""
|
||||
|
||||
engine: Engine
|
||||
session: scoped_session[Session]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def sqlite_database(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sqlite_engine: Engine,
|
||||
) -> Iterator[scoped_session[Session]]:
|
||||
"""Route controller transactions and model author lookups through SQLite."""
|
||||
|
||||
database_session = scoped_session(sessionmaker(bind=sqlite_engine, expire_on_commit=False))
|
||||
database = SQLiteDatabase(engine=sqlite_engine, session=database_session)
|
||||
monkeypatch.setattr(workflow_controller, "db", database)
|
||||
monkeypatch.setattr(workflow_models, "db", database)
|
||||
|
||||
with database_session() as session:
|
||||
default_author = Account(name="Default Author", email="default-author@example.com")
|
||||
default_author.id = DEFAULT_WORKFLOW_CREATED_BY
|
||||
session.add(default_author)
|
||||
session.commit()
|
||||
|
||||
try:
|
||||
yield database_session
|
||||
finally:
|
||||
database_session.remove()
|
||||
|
||||
|
||||
def empty_mapping() -> dict[str, object]:
|
||||
return {}
|
||||
|
||||
@@ -206,18 +237,15 @@ def make_pipeline(
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def workflow_author(db_session_with_containers: Session) -> Account:
|
||||
def workflow_author(sqlite_database: scoped_session[Session]) -> Account:
|
||||
account = Account(name="Alice", email=f"alice-{uuid4()}@example.com")
|
||||
db_session_with_containers.add(account)
|
||||
db_session_with_containers.commit()
|
||||
account.id = str(uuid4())
|
||||
sqlite_database.add(account)
|
||||
sqlite_database.commit()
|
||||
return account
|
||||
|
||||
|
||||
class TestDraftWorkflowApi:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask) -> Flask:
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_get_draft_success(self, app: Flask, workflow_author: Account) -> None:
|
||||
api = DraftRagPipelineApi()
|
||||
method = unwrap(api.get)
|
||||
@@ -278,7 +306,6 @@ class TestDraftWorkflowApi:
|
||||
|
||||
with (
|
||||
app.test_request_context("/", json={"graph": empty_mapping(), "features": empty_mapping()}),
|
||||
patch.object(type(console_ns), "payload", {"graph": empty_mapping(), "features": empty_mapping()}),
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService",
|
||||
return_value=service,
|
||||
@@ -373,10 +400,6 @@ class TestDraftWorkflowApi:
|
||||
|
||||
|
||||
class TestDraftRunNodes:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask) -> Flask:
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_iteration_node_success(self, app: Flask) -> None:
|
||||
api = RagPipelineDraftRunIterationNodeApi()
|
||||
method = unwrap(api.post)
|
||||
@@ -386,7 +409,6 @@ class TestDraftRunNodes:
|
||||
|
||||
with (
|
||||
app.test_request_context("/", json={"inputs": empty_mapping()}),
|
||||
patch.object(type(console_ns), "payload", {"inputs": empty_mapping()}),
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService.generate_single_iteration",
|
||||
return_value=MagicMock(),
|
||||
@@ -408,7 +430,6 @@ class TestDraftRunNodes:
|
||||
|
||||
with (
|
||||
app.test_request_context("/", json={"inputs": empty_mapping()}),
|
||||
patch.object(type(console_ns), "payload", {"inputs": empty_mapping()}),
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService.generate_single_iteration",
|
||||
side_effect=services.errors.conversation.ConversationNotExistsError(),
|
||||
@@ -426,7 +447,6 @@ class TestDraftRunNodes:
|
||||
|
||||
with (
|
||||
app.test_request_context("/", json={"inputs": empty_mapping()}),
|
||||
patch.object(type(console_ns), "payload", {"inputs": empty_mapping()}),
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService.generate_single_loop",
|
||||
return_value=MagicMock(),
|
||||
@@ -439,86 +459,7 @@ class TestDraftRunNodes:
|
||||
assert method(api, user, pipeline, "node") == {"ok": True}
|
||||
|
||||
|
||||
class TestPipelineRunApis:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask) -> Flask:
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_draft_run_success(self, app: Flask) -> None:
|
||||
api = DraftRagPipelineRunApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
pipeline = make_pipeline()
|
||||
user = make_account()
|
||||
session = MagicMock(spec=Session)
|
||||
|
||||
payload = {
|
||||
"inputs": empty_mapping(),
|
||||
"datasource_type": "x",
|
||||
"datasource_info_list": empty_list(),
|
||||
"start_node_id": "n",
|
||||
}
|
||||
|
||||
with (
|
||||
app.test_request_context("/", json=payload),
|
||||
patch.object(type(console_ns), "payload", payload),
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.load_rag_pipeline",
|
||||
return_value=pipeline,
|
||||
) as load_pipeline,
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService.generate",
|
||||
return_value=MagicMock(),
|
||||
) as generate,
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.helper.compact_generate_response",
|
||||
return_value={"ok": True},
|
||||
),
|
||||
):
|
||||
assert method(api, session, user, pipeline.id) == {"ok": True}
|
||||
|
||||
load_pipeline.assert_called_once_with(session, pipeline.id)
|
||||
assert generate.call_args.kwargs["session"] is session
|
||||
|
||||
def test_draft_run_rate_limit(self, app: Flask) -> None:
|
||||
api = DraftRagPipelineRunApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
pipeline = make_pipeline()
|
||||
user = make_account()
|
||||
session = MagicMock(spec=Session)
|
||||
payload: dict[str, object] = {
|
||||
"inputs": empty_mapping(),
|
||||
"datasource_type": "x",
|
||||
"datasource_info_list": empty_list(),
|
||||
"start_node_id": "n",
|
||||
}
|
||||
|
||||
with (
|
||||
app.test_request_context("/", json=payload),
|
||||
patch.object(
|
||||
type(console_ns),
|
||||
"payload",
|
||||
payload,
|
||||
),
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.load_rag_pipeline",
|
||||
return_value=pipeline,
|
||||
),
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService.generate",
|
||||
side_effect=InvokeRateLimitError("limit"),
|
||||
),
|
||||
):
|
||||
with pytest.raises(InvokeRateLimitHttpError):
|
||||
method(api, session, user, pipeline.id)
|
||||
|
||||
|
||||
class TestDraftNodeRun:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask) -> Flask:
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_execution_not_found(self, app: Flask) -> None:
|
||||
api = RagPipelineDraftNodeRunApi()
|
||||
method = unwrap(api.post)
|
||||
@@ -531,7 +472,6 @@ class TestDraftNodeRun:
|
||||
|
||||
with (
|
||||
app.test_request_context("/", json={"inputs": empty_mapping()}),
|
||||
patch.object(type(console_ns), "payload", {"inputs": empty_mapping()}),
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService",
|
||||
return_value=service,
|
||||
@@ -542,11 +482,7 @@ class TestDraftNodeRun:
|
||||
|
||||
|
||||
class TestPublishedPipelineApis:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask) -> Flask:
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_publish_success(self, app: Flask, db_session_with_containers: Session) -> None:
|
||||
def test_publish_success(self, app: Flask) -> None:
|
||||
api = PublishedRagPipelineApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
@@ -557,10 +493,6 @@ class TestPublishedPipelineApis:
|
||||
description="test",
|
||||
created_by=str(uuid4()),
|
||||
)
|
||||
db_session_with_containers.add(pipeline)
|
||||
db_session_with_containers.commit()
|
||||
db_session_with_containers.expire_all()
|
||||
|
||||
user = make_account(id="u1")
|
||||
|
||||
workflow = make_workflow(id=str(uuid4()), created_at=naive_utc_now())
|
||||
@@ -582,10 +514,6 @@ class TestPublishedPipelineApis:
|
||||
|
||||
|
||||
class TestMiscApis:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask) -> Flask:
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_task_stop(self, app: Flask) -> None:
|
||||
api = RagPipelineTaskStopApi()
|
||||
method = unwrap(api.post)
|
||||
@@ -603,18 +531,6 @@ class TestMiscApis:
|
||||
stop_mock.assert_called_once()
|
||||
assert result["result"] == "success"
|
||||
|
||||
def test_transform_forbidden(self, app: Flask) -> None:
|
||||
api = RagPipelineTransformApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
user = make_account(role=TenantAccountRole.NORMAL)
|
||||
|
||||
with (
|
||||
app.test_request_context("/"),
|
||||
):
|
||||
with pytest.raises(Forbidden):
|
||||
method(api, MagicMock(spec=Session), user, "ds1")
|
||||
|
||||
def test_recommended_plugins(self, app: Flask) -> None:
|
||||
api = RagPipelineRecommendedPluginApi()
|
||||
method = unwrap(api.get)
|
||||
@@ -640,85 +556,7 @@ class TestMiscApis:
|
||||
service.get_recommended_plugins.assert_called_once_with("all", user, tenant_id)
|
||||
|
||||
|
||||
class TestPublishedRagPipelineRunApi:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask) -> Flask:
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_published_run_success(self, app: Flask) -> None:
|
||||
api = PublishedRagPipelineRunApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
pipeline = make_pipeline()
|
||||
user = make_account()
|
||||
session = MagicMock(spec=Session)
|
||||
|
||||
payload = {
|
||||
"inputs": empty_mapping(),
|
||||
"datasource_type": "x",
|
||||
"datasource_info_list": empty_list(),
|
||||
"start_node_id": "n",
|
||||
"response_mode": "blocking",
|
||||
}
|
||||
|
||||
with (
|
||||
app.test_request_context("/", json=payload),
|
||||
patch.object(type(console_ns), "payload", payload),
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.load_rag_pipeline",
|
||||
return_value=pipeline,
|
||||
) as load_pipeline,
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService.generate",
|
||||
return_value=MagicMock(),
|
||||
) as generate,
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.helper.compact_generate_response",
|
||||
return_value={"ok": True},
|
||||
),
|
||||
):
|
||||
result = method(api, session, user, pipeline.id)
|
||||
assert result == {"ok": True}
|
||||
|
||||
load_pipeline.assert_called_once_with(session, pipeline.id)
|
||||
assert generate.call_args.kwargs["session"] is session
|
||||
|
||||
def test_published_run_rate_limit(self, app: Flask) -> None:
|
||||
api = PublishedRagPipelineRunApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
pipeline = make_pipeline()
|
||||
user = make_account()
|
||||
session = MagicMock(spec=Session)
|
||||
|
||||
payload = {
|
||||
"inputs": empty_mapping(),
|
||||
"datasource_type": "x",
|
||||
"datasource_info_list": empty_list(),
|
||||
"start_node_id": "n",
|
||||
}
|
||||
|
||||
with (
|
||||
app.test_request_context("/", json=payload),
|
||||
patch.object(type(console_ns), "payload", payload),
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.load_rag_pipeline",
|
||||
return_value=pipeline,
|
||||
),
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.PipelineGenerateService.generate",
|
||||
side_effect=InvokeRateLimitError("limit"),
|
||||
),
|
||||
):
|
||||
with pytest.raises(InvokeRateLimitHttpError):
|
||||
method(api, session, user, pipeline.id)
|
||||
|
||||
|
||||
class TestDefaultBlockConfigApi:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask) -> Flask:
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_get_block_config_success(self, app: Flask) -> None:
|
||||
api = DefaultRagPipelineBlockConfigApi()
|
||||
method = unwrap(api.get)
|
||||
@@ -750,10 +588,6 @@ class TestDefaultBlockConfigApi:
|
||||
|
||||
|
||||
class TestPublishedAllRagPipelineApi:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask) -> Flask:
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_get_published_workflows_success(self, app: Flask) -> None:
|
||||
api = PublishedAllRagPipelineApi()
|
||||
method = unwrap(api.get)
|
||||
@@ -792,10 +626,6 @@ class TestPublishedAllRagPipelineApi:
|
||||
|
||||
|
||||
class TestRagPipelineByIdApi:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask) -> Flask:
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_patch_success(self, app: Flask) -> None:
|
||||
api = RagPipelineByIdApi()
|
||||
method = unwrap(api.patch)
|
||||
@@ -812,7 +642,6 @@ class TestRagPipelineByIdApi:
|
||||
|
||||
with (
|
||||
app.test_request_context("/", json=payload),
|
||||
patch.object(type(console_ns), "payload", payload),
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService",
|
||||
return_value=service,
|
||||
@@ -831,10 +660,7 @@ class TestRagPipelineByIdApi:
|
||||
pipeline = make_pipeline()
|
||||
user = make_account()
|
||||
|
||||
with (
|
||||
app.test_request_context("/", json={}),
|
||||
patch.object(type(console_ns), "payload", empty_mapping()),
|
||||
):
|
||||
with app.test_request_context("/", json={}):
|
||||
result, status = method(api, user, pipeline, "w1")
|
||||
assert status == 400
|
||||
|
||||
@@ -870,10 +696,6 @@ class TestRagPipelineByIdApi:
|
||||
|
||||
|
||||
class TestRagPipelineWorkflowLastRunApi:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask) -> Flask:
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_last_run_success(self, app: Flask) -> None:
|
||||
api = RagPipelineWorkflowLastRunApi()
|
||||
method = unwrap(api.get)
|
||||
@@ -919,10 +741,6 @@ class TestRagPipelineWorkflowLastRunApi:
|
||||
|
||||
|
||||
class TestRagPipelineWorkflowRunNodeExecutionListApi:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask) -> Flask:
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_get_node_executions_passes_current_user(self, app: Flask) -> None:
|
||||
api = RagPipelineWorkflowRunNodeExecutionListApi()
|
||||
method = unwrap(api.get)
|
||||
@@ -955,10 +773,6 @@ class TestRagPipelineWorkflowRunNodeExecutionListApi:
|
||||
|
||||
|
||||
class TestRagPipelineDatasourceVariableApi:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask) -> Flask:
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_set_datasource_variables_success(self, app: Flask) -> None:
|
||||
api = RagPipelineDatasourceVariableApi()
|
||||
method = unwrap(api.post)
|
||||
@@ -978,7 +792,6 @@ class TestRagPipelineDatasourceVariableApi:
|
||||
|
||||
with (
|
||||
app.test_request_context("/", json=payload),
|
||||
patch.object(type(console_ns), "payload", payload),
|
||||
patch(
|
||||
"controllers.console.datasets.rag_pipeline.rag_pipeline_workflow.RagPipelineService",
|
||||
return_value=service,
|
||||
Reference in New Issue
Block a user