refactor(api): migrate console.datasets.data_source to BaseModel (#36624)

This commit is contained in:
chariri
2026-06-03 19:38:39 +00:00
committed by GitHub
parent 4fc62d3b38
commit d3058d63bd
14 changed files with 682 additions and 419 deletions
@@ -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"] == {}