mirror of
https://github.com/langgenius/dify.git
synced 2026-09-24 23:22:26 +08:00
refactor(api): migrate console.datasets.data_source to BaseModel (#36624)
This commit is contained in:
+57
-10
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import MagicMock, PropertyMock, patch
|
||||
|
||||
import pytest
|
||||
@@ -11,12 +12,14 @@ from werkzeug.exceptions import NotFound
|
||||
from controllers.console.datasets import data_source
|
||||
from controllers.console.datasets.data_source import (
|
||||
DataSourceApi,
|
||||
DataSourceNotionApi,
|
||||
DataSourceNotionDatasetSyncApi,
|
||||
DataSourceNotionDocumentSyncApi,
|
||||
DataSourceNotionIndexingEstimateApi,
|
||||
DataSourceNotionListApi,
|
||||
DataSourceNotionPreviewApi,
|
||||
)
|
||||
from core.rag.index_processor.constant.index_type import IndexStructureType
|
||||
from models import DataSourceOauthBinding
|
||||
|
||||
|
||||
def unwrap(func):
|
||||
@@ -59,13 +62,29 @@ class TestDataSourceApi:
|
||||
api = DataSourceApi()
|
||||
method = unwrap(api.get)
|
||||
|
||||
binding = MagicMock(
|
||||
id="b1",
|
||||
binding = DataSourceOauthBinding(
|
||||
tenant_id="tenant-1",
|
||||
access_token="token",
|
||||
provider="notion",
|
||||
created_at="now",
|
||||
disabled=False,
|
||||
source_info={},
|
||||
source_info={
|
||||
"workspace_name": "Workspace",
|
||||
"workspace_id": "workspace-1",
|
||||
"workspace_icon": None,
|
||||
"total": 1,
|
||||
"pages": [
|
||||
{
|
||||
"page_id": "page-1",
|
||||
"page_name": "Page",
|
||||
"page_icon": {"type": "emoji", "emoji": "P", "url": None},
|
||||
"parent_id": "parent-1",
|
||||
"type": "page",
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
binding.id = "b1"
|
||||
binding.created_at = datetime(2026, 5, 25, 1, 2, 3, tzinfo=UTC)
|
||||
binding.disabled = False
|
||||
|
||||
with (
|
||||
app.test_request_context("/"),
|
||||
@@ -77,7 +96,29 @@ class TestDataSourceApi:
|
||||
response, status = method(api)
|
||||
|
||||
assert status == 200
|
||||
assert response["data"][0]["is_bound"] is True
|
||||
assert response["data"][0] == {
|
||||
"id": "b1",
|
||||
"provider": "notion",
|
||||
"created_at": 1779670923,
|
||||
"is_bound": True,
|
||||
"disabled": False,
|
||||
"source_info": {
|
||||
"workspace_name": "Workspace",
|
||||
"workspace_id": "workspace-1",
|
||||
"workspace_icon": None,
|
||||
"pages": [
|
||||
{
|
||||
"page_name": "Page",
|
||||
"page_id": "page-1",
|
||||
"page_icon": {"type": "emoji", "url": None, "emoji": "P"},
|
||||
"parent_id": "parent-1",
|
||||
"type": "page",
|
||||
}
|
||||
],
|
||||
"total": 1,
|
||||
},
|
||||
"link": "http://localhost/console/api/oauth/data-source/notion",
|
||||
}
|
||||
|
||||
def test_get_no_bindings(self, app: Flask, patch_tenant):
|
||||
api = DataSourceApi()
|
||||
@@ -322,13 +363,13 @@ class TestDataSourceNotionListApi:
|
||||
method(api)
|
||||
|
||||
|
||||
class TestDataSourceNotionApi:
|
||||
class TestDataSourceNotionPreviewApi:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask):
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_get_preview_success(self, app: Flask, patch_tenant):
|
||||
api = DataSourceNotionApi()
|
||||
api = DataSourceNotionPreviewApi()
|
||||
method = unwrap(api.get)
|
||||
|
||||
extractor = MagicMock(extract=lambda: [MagicMock(page_content="hello")])
|
||||
@@ -348,8 +389,14 @@ class TestDataSourceNotionApi:
|
||||
|
||||
assert status == 200
|
||||
|
||||
|
||||
class TestDataSourceNotionIndexingEstimateApi:
|
||||
@pytest.fixture
|
||||
def app(self, flask_app_with_containers: Flask):
|
||||
return flask_app_with_containers
|
||||
|
||||
def test_post_indexing_estimate_success(self, app: Flask, patch_tenant):
|
||||
api = DataSourceNotionApi()
|
||||
api = DataSourceNotionIndexingEstimateApi()
|
||||
method = unwrap(api.post)
|
||||
|
||||
payload = {
|
||||
|
||||
@@ -0,0 +1,173 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
from datetime import UTC, datetime
|
||||
from typing import cast
|
||||
from unittest.mock import MagicMock, PropertyMock, patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
|
||||
from controllers.console.datasets import data_source as module
|
||||
from controllers.console.datasets.data_source import DataSourceApi, DataSourceNotionListApi
|
||||
from models import DataSourceOauthBinding
|
||||
|
||||
ControllerMethod = Callable[..., tuple[dict[str, object], int]]
|
||||
|
||||
|
||||
def unwrap(func: object) -> ControllerMethod:
|
||||
return cast(ControllerMethod, inspect.unwrap(cast(Callable[..., object], func)))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def flask_app() -> Flask:
|
||||
app = Flask(__name__)
|
||||
app.config["TESTING"] = True
|
||||
return app
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tenant_context() -> tuple[MagicMock, str]:
|
||||
return MagicMock(id="user-1"), "tenant-1"
|
||||
|
||||
|
||||
def test_get_data_source_integrates_serializes_orm_binding(
|
||||
flask_app: Flask, tenant_context: tuple[MagicMock, str]
|
||||
) -> None:
|
||||
binding = DataSourceOauthBinding(
|
||||
tenant_id="tenant-1",
|
||||
access_token="token",
|
||||
provider="notion",
|
||||
source_info={
|
||||
"workspace_name": "Workspace",
|
||||
"workspace_id": "workspace-1",
|
||||
"workspace_icon": None,
|
||||
"total": 1,
|
||||
"pages": [
|
||||
{
|
||||
"page_id": "page-1",
|
||||
"page_name": "Page",
|
||||
"page_icon": {"type": "emoji", "emoji": "P", "url": None},
|
||||
"parent_id": "parent-1",
|
||||
"type": "page",
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
binding.id = "binding-1"
|
||||
binding.created_at = datetime(2026, 5, 25, 1, 2, 3, tzinfo=UTC)
|
||||
binding.disabled = False
|
||||
|
||||
with (
|
||||
flask_app.test_request_context("/"),
|
||||
patch.object(module, "current_account_with_tenant", return_value=tenant_context),
|
||||
patch.object(module.db.session, "scalars", return_value=MagicMock(all=lambda: [binding])),
|
||||
):
|
||||
response, status = unwrap(DataSourceApi().get)(DataSourceApi())
|
||||
|
||||
assert status == 200
|
||||
assert response == {
|
||||
"data": [
|
||||
{
|
||||
"id": "binding-1",
|
||||
"provider": "notion",
|
||||
"created_at": 1779670923,
|
||||
"is_bound": True,
|
||||
"disabled": False,
|
||||
"source_info": {
|
||||
"workspace_name": "Workspace",
|
||||
"workspace_id": "workspace-1",
|
||||
"workspace_icon": None,
|
||||
"pages": [
|
||||
{
|
||||
"page_name": "Page",
|
||||
"page_id": "page-1",
|
||||
"page_icon": {"type": "emoji", "url": None, "emoji": "P"},
|
||||
"parent_id": "parent-1",
|
||||
"type": "page",
|
||||
}
|
||||
],
|
||||
"total": 1,
|
||||
},
|
||||
"link": "http://localhost/console/api/oauth/data-source/notion",
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def test_get_data_source_integrates_preserves_empty_list_when_no_binding(
|
||||
flask_app: Flask, tenant_context: tuple[MagicMock, str]
|
||||
) -> None:
|
||||
with (
|
||||
flask_app.test_request_context("/"),
|
||||
patch.object(module, "current_account_with_tenant", return_value=tenant_context),
|
||||
patch.object(module.db.session, "scalars", return_value=MagicMock(all=lambda: [])),
|
||||
):
|
||||
response, status = unwrap(DataSourceApi().get)(DataSourceApi())
|
||||
|
||||
assert status == 200
|
||||
assert response == {"data": []}
|
||||
|
||||
|
||||
def test_notion_pre_import_pages_serializes_frontend_list_shape(
|
||||
flask_app: Flask, tenant_context: tuple[MagicMock, str]
|
||||
) -> None:
|
||||
page = MagicMock(
|
||||
page_id="page-1",
|
||||
page_name="Page",
|
||||
type="page",
|
||||
parent_id="parent-1",
|
||||
page_icon={"type": "emoji", "emoji": "P", "url": None},
|
||||
)
|
||||
online_document_message = MagicMock(
|
||||
result=[
|
||||
MagicMock(
|
||||
workspace_id="workspace-1",
|
||||
workspace_name="Workspace",
|
||||
workspace_icon=None,
|
||||
pages=[page],
|
||||
)
|
||||
]
|
||||
)
|
||||
runtime = MagicMock(
|
||||
get_online_document_pages=MagicMock(return_value=iter([online_document_message])),
|
||||
datasource_provider_type=MagicMock(return_value="online_document"),
|
||||
)
|
||||
|
||||
with (
|
||||
flask_app.test_request_context("/?credential_id=credential-1"),
|
||||
patch.object(module, "current_account_with_tenant", return_value=tenant_context),
|
||||
patch.object(
|
||||
module.DatasourceProviderService,
|
||||
"get_datasource_credentials",
|
||||
return_value={"token": "token"},
|
||||
),
|
||||
patch.object(type(module.db), "engine", new_callable=PropertyMock, return_value=MagicMock()),
|
||||
patch.object(module, "sessionmaker"),
|
||||
patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime),
|
||||
):
|
||||
response, status = unwrap(DataSourceNotionListApi().get)(DataSourceNotionListApi())
|
||||
|
||||
assert status == 200
|
||||
assert response == {
|
||||
"notion_info": [
|
||||
{
|
||||
"workspace_name": "Workspace",
|
||||
"workspace_id": "workspace-1",
|
||||
"workspace_icon": None,
|
||||
"pages": [
|
||||
{
|
||||
"page_name": "Page",
|
||||
"page_id": "page-1",
|
||||
"page_icon": {"type": "emoji", "url": None, "emoji": "P"},
|
||||
"parent_id": "parent-1",
|
||||
"type": "page",
|
||||
"is_bound": False,
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
runtime.get_online_document_pages.assert_called_once()
|
||||
assert runtime.get_online_document_pages.call_args.kwargs["datasource_parameters"] == {}
|
||||
Reference in New Issue
Block a user