opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-28 15:49:42 +08:00
Code Issues Packages Projects Releases Wiki Activity

test: move data source controller coverage to unit tests (#38924)

Browse Source
This commit is contained in:
Asuka Minato 2026-07-28 12:51:22 +09:00 committed by GitHub
parent 3c7ad816d9
commit e1ce808567
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 371 additions and 515 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

551
api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py
View File

@ -1,494 +1,85 @@
"""Testcontainers integration tests for controllers.console.datasets.data_source endpoints."""
"""Integration coverage for Notion page bindings backed by persisted documents."""
from __future__ import annotations
import inspect
from collections.abc import Iterator
from datetime import UTC, datetime
from unittest.mock import MagicMock, PropertyMock, patch
from inspect import unwrap
from unittest.mock import MagicMock, patch
from uuid import uuid4
import pytest
from flask import Flask
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from controllers.console.datasets import data_source
from controllers.console.datasets.data_source import (
DataSourceApi,
DataSourceNotionDatasetSyncApi,
DataSourceNotionDocumentSyncApi,
DataSourceNotionIndexingEstimateApi,
DataSourceNotionListApi,
DataSourceNotionPreviewApi,
)
from core.rag.index_processor.constant.index_type import IndexStructureType
from models import Account, DataSourceOauthBinding
from controllers.console.datasets.data_source import DataSourceNotionListApi
from models import Account
from models.dataset import Document
from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus
@pytest.fixture
def current_user() -> Account:
account = Account(name="Test User", email="u1@example.com")
account.id = "u1"
return account
def test_notion_page_is_marked_bound_from_persisted_document(
flask_app_with_containers: Flask,
db_session_with_containers: Session,
) -> None:
tenant_id = str(uuid4())
dataset_id = str(uuid4())
account = Account(name="Test User", email="user@example.com")
account.id = str(uuid4())
document = Document(
tenant_id=tenant_id,
dataset_id=dataset_id,
position=1,
data_source_type=DataSourceType.NOTION_IMPORT,
data_source_info='{"notion_page_id": "page-1"}',
batch=f"batch-{uuid4()}",
name="Notion Page",
created_from=DocumentCreatedFrom.WEB,
created_by=str(uuid4()),
indexing_status=IndexingStatus.COMPLETED,
enabled=True,
)
db_session_with_containers.add(document)
db_session_with_containers.commit()
runtime = MagicMock(
get_online_document_pages=lambda **_kwargs: iter(
[
MagicMock(
result=[
MagicMock(
workspace_id="workspace-1",
workspace_name="Workspace",
workspace_icon=None,
pages=[
MagicMock(
page_id="page-1",
page_name="Page",
type="page",
parent_id="parent",
page_icon=None,
)
],
)
]
)
]
),
datasource_provider_type=lambda: None,
)
@pytest.fixture
def mock_engine() -> Iterator[None]:
with patch.object(
type(data_source.db),
"engine",
new_callable=PropertyMock,
return_value=MagicMock(),
with (
flask_app_with_containers.test_request_context(f"/?credential_id=c1&dataset_id={dataset_id}"),
patch(
"controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials",
return_value={"token": "token"},
),
patch(
"controllers.console.datasets.data_source.DatasetService.get_dataset",
return_value=MagicMock(data_source_type="notion_import"),
),
patch(
"core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime",
return_value=runtime,
),
):
yield
class TestDataSourceApi:
@pytest.fixture
def app(self, flask_app_with_containers: Flask) -> Flask:
return flask_app_with_containers
def test_get_success(self, app: Flask) -> None:
api = DataSourceApi()
method = inspect.unwrap(api.get)
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 = "b1"
binding.created_at = datetime(2026, 5, 25, 1, 2, 3, tzinfo=UTC)
binding.disabled = False
with (
app.test_request_context("/"),
patch(
"controllers.console.datasets.data_source.db.session.scalars",
return_value=MagicMock(all=lambda: [binding]),
),
):
response, status = method(api, "tenant-1")
assert status == 200
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) -> None:
api = DataSourceApi()
method = inspect.unwrap(api.get)
with (
app.test_request_context("/"),
patch(
"controllers.console.datasets.data_source.db.session.scalars",
return_value=MagicMock(all=lambda: []),
),
):
response, status = method(api, "tenant-1")
assert status == 200
assert response["data"] == []
def test_patch_enable_binding(self, app: Flask) -> None:
api = DataSourceApi()
method = inspect.unwrap(api.patch)
binding = MagicMock(id="b1", disabled=True)
session = MagicMock()
session.scalar.return_value = binding
with app.test_request_context("/"):
response, status = method(api, session, "tenant-1", "b1", "enable")
assert status == 200
assert binding.disabled is False
def test_patch_disable_binding(self, app: Flask) -> None:
api = DataSourceApi()
method = inspect.unwrap(api.patch)
binding = MagicMock(id="b1", disabled=False)
session = MagicMock()
session.scalar.return_value = binding
with app.test_request_context("/"):
response, status = method(api, session, "tenant-1", "b1", "disable")
assert status == 200
assert binding.disabled is True
def test_patch_binding_not_found(self, app: Flask) -> None:
api = DataSourceApi()
method = inspect.unwrap(api.patch)
session = MagicMock()
session.scalar.return_value = None
with app.test_request_context("/"):
with pytest.raises(NotFound):
method(api, session, "tenant-1", "b1", "enable")
def test_patch_enable_already_enabled(self, app: Flask) -> None:
api = DataSourceApi()
method = inspect.unwrap(api.patch)
binding = MagicMock(id="b1", disabled=False)
session = MagicMock()
session.scalar.return_value = binding
with app.test_request_context("/"):
with pytest.raises(ValueError):
method(api, session, "tenant-1", "b1", "enable")
def test_patch_disable_already_disabled(self, app: Flask) -> None:
api = DataSourceApi()
method = inspect.unwrap(api.patch)
binding = MagicMock(id="b1", disabled=True)
session = MagicMock()
session.scalar.return_value = binding
with app.test_request_context("/"):
with pytest.raises(ValueError):
method(api, session, "tenant-1", "b1", "disable")
class TestDataSourceNotionListApi:
@pytest.fixture
def app(self, flask_app_with_containers: Flask) -> Flask:
return flask_app_with_containers
def test_get_credential_not_found(self, app: Flask, current_user: Account) -> None:
api = DataSourceNotionListApi()
method = inspect.unwrap(api.get)
with (
app.test_request_context("/?credential_id=c1"),
patch(
"controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials",
return_value=None,
),
):
with pytest.raises(NotFound):
method(api, MagicMock(), "tenant-1", current_user)
def test_get_success_no_dataset_id(self, app: Flask, current_user: Account, mock_engine: None) -> None:
api = DataSourceNotionListApi()
method = inspect.unwrap(api.get)
page = MagicMock(
page_id="p1",
page_name="Page 1",
type="page",
parent_id="parent",
page_icon=None,
response, status = unwrap(DataSourceNotionListApi().get)(
DataSourceNotionListApi(), db_session_with_containers, tenant_id, account
)
online_document_message = MagicMock(
result=[
MagicMock(
workspace_id="w1",
workspace_name="My Workspace",
workspace_icon="icon",
pages=[page],
)
]
)
with (
app.test_request_context("/?credential_id=c1"),
patch(
"controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials",
return_value={"token": "t"},
),
patch(
"core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime",
return_value=MagicMock(
get_online_document_pages=lambda **kw: iter([online_document_message]),
datasource_provider_type=lambda: None,
),
),
):
response, status = method(api, MagicMock(), "tenant-1", current_user)
assert status == 200
def test_get_success_with_dataset_id(
self, app: Flask, current_user: Account, mock_engine: None, db_session_with_containers: Session
) -> None:
api = DataSourceNotionListApi()
method = inspect.unwrap(api.get)
tenant_id = str(uuid4())
dataset_id = str(uuid4())
page = MagicMock(
page_id="p1",
page_name="Page 1",
type="page",
parent_id="parent",
page_icon=None,
)
online_document_message = MagicMock(
result=[
MagicMock(
workspace_id="w1",
workspace_name="My Workspace",
workspace_icon="icon",
pages=[page],
)
]
)
dataset = MagicMock(data_source_type="notion_import")
document = Document(
tenant_id=tenant_id,
dataset_id=dataset_id,
position=1,
data_source_type=DataSourceType.NOTION_IMPORT,
data_source_info='{"notion_page_id": "p1"}',
batch=f"batch-{uuid4()}",
name="Notion Page",
created_from=DocumentCreatedFrom.WEB,
created_by=str(uuid4()),
indexing_status=IndexingStatus.COMPLETED,
enabled=True,
)
db_session_with_containers.add(document)
db_session_with_containers.commit()
with (
app.test_request_context(f"/?credential_id=c1&dataset_id={dataset_id}"),
patch(
"controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials",
return_value={"token": "t"},
),
patch(
"controllers.console.datasets.data_source.DatasetService.get_dataset",
return_value=dataset,
),
patch(
"core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime",
return_value=MagicMock(
get_online_document_pages=lambda **kw: iter([online_document_message]),
datasource_provider_type=lambda: None,
),
),
):
response, status = method(api, db_session_with_containers, tenant_id, current_user)
assert status == 200
def test_get_invalid_dataset_type(self, app: Flask, current_user: Account, mock_engine: None) -> None:
api = DataSourceNotionListApi()
method = inspect.unwrap(api.get)
dataset = MagicMock(data_source_type="other_type")
with (
app.test_request_context("/?credential_id=c1&dataset_id=ds1"),
patch(
"controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials",
return_value={"token": "t"},
),
patch(
"controllers.console.datasets.data_source.DatasetService.get_dataset",
return_value=dataset,
),
):
with pytest.raises(ValueError):
method(api, MagicMock(), "tenant-1", current_user)
class TestDataSourceNotionPreviewApi:
@pytest.fixture
def app(self, flask_app_with_containers: Flask) -> Flask:
return flask_app_with_containers
def test_get_preview_success(self, app: Flask) -> None:
api = DataSourceNotionPreviewApi()
method = inspect.unwrap(api.get)
extractor = MagicMock(extract=lambda: [MagicMock(page_content="hello")])
with (
app.test_request_context("/?credential_id=c1"),
patch(
"controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials",
return_value={"integration_secret": "t"},
),
patch(
"controllers.console.datasets.data_source.NotionExtractor",
return_value=extractor,
),
):
response, status = method(api, "tenant-1", "p1", "page")
assert status == 200
class TestDataSourceNotionIndexingEstimateApi:
@pytest.fixture
def app(self, flask_app_with_containers: Flask) -> Flask:
return flask_app_with_containers
def test_post_indexing_estimate_success(self, app: Flask) -> None:
api = DataSourceNotionIndexingEstimateApi()
method = inspect.unwrap(api.post)
empty_rules: dict[str, object] = {}
payload: dict[str, object] = {
"notion_info_list": [
{
"workspace_id": "w1",
"credential_id": "c1",
"pages": [{"page_id": "p1", "type": "page"}],
}
],
"process_rule": {"rules": empty_rules},
"doc_form": IndexStructureType.PARAGRAPH_INDEX,
"doc_language": "English",
}
with (
app.test_request_context("/", method="POST", json=payload, headers={"Content-Type": "application/json"}),
patch(
"controllers.console.datasets.data_source.DocumentService.estimate_args_validate",
),
patch(
"controllers.console.datasets.data_source.IndexingRunner.indexing_estimate",
return_value=MagicMock(model_dump=lambda: {"total_pages": 1}),
),
):
response, status = method(api, MagicMock(), "tenant-1")
assert status == 200
class TestDataSourceNotionDatasetSyncApi:
@pytest.fixture
def app(self, flask_app_with_containers: Flask) -> Flask:
return flask_app_with_containers
def test_get_success(self, app: Flask) -> None:
api = DataSourceNotionDatasetSyncApi()
method = inspect.unwrap(api.get)
with (
app.test_request_context("/"),
patch(
"controllers.console.datasets.data_source.DatasetService.get_dataset",
return_value=MagicMock(),
),
patch(
"controllers.console.datasets.data_source.DocumentService.get_document_by_dataset_id",
return_value=[MagicMock(id="d1")],
),
patch(
"controllers.console.datasets.data_source.document_indexing_sync_task.delay",
return_value=None,
),
):
response, status = method(api, MagicMock(), "ds-1")
assert status == 200
def test_get_dataset_not_found(self, app: Flask) -> None:
api = DataSourceNotionDatasetSyncApi()
method = inspect.unwrap(api.get)
with (
app.test_request_context("/"),
patch(
"controllers.console.datasets.data_source.DatasetService.get_dataset",
return_value=None,
),
):
with pytest.raises(NotFound):
method(api, MagicMock(), "ds-1")
class TestDataSourceNotionDocumentSyncApi:
@pytest.fixture
def app(self, flask_app_with_containers: Flask) -> Flask:
return flask_app_with_containers
def test_get_success(self, app: Flask) -> None:
api = DataSourceNotionDocumentSyncApi()
method = inspect.unwrap(api.get)
with (
app.test_request_context("/"),
patch(
"controllers.console.datasets.data_source.DatasetService.get_dataset",
return_value=MagicMock(),
),
patch(
"controllers.console.datasets.data_source.DocumentService.get_document",
return_value=MagicMock(),
),
patch(
"controllers.console.datasets.data_source.document_indexing_sync_task.delay",
return_value=None,
),
):
response, status = method(api, MagicMock(), "ds-1", "doc-1")
assert status == 200
def test_get_document_not_found(self, app: Flask) -> None:
api = DataSourceNotionDocumentSyncApi()
method = inspect.unwrap(api.get)
with (
app.test_request_context("/"),
patch(
"controllers.console.datasets.data_source.DatasetService.get_dataset",
return_value=MagicMock(),
),
patch(
"controllers.console.datasets.data_source.DocumentService.get_document",
return_value=None,
),
):
with pytest.raises(NotFound):
method(api, MagicMock(), "ds-1", "doc-1")
assert status == 200
assert response["notion_info"][0]["pages"][0]["is_bound"] is True

164
api/tests/unit_tests/controllers/console/datasets/test_data_source.py
View File

@ -1,18 +1,22 @@
from __future__ import annotations
import inspect
from collections.abc import Callable
from collections.abc import Callable, Iterator
from datetime import UTC, datetime
from typing import cast
from typing import Literal, cast
from unittest.mock import MagicMock, PropertyMock, patch
from uuid import uuid4
from uuid import UUID
import pytest
from flask import Flask
from sqlalchemy import select
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from controllers.console.datasets import data_source as module
from controllers.console.datasets.data_source import DataSourceApi, DataSourceNotionListApi
from models import Account, DataSourceOauthBinding
from models.engine import db
ControllerMethod = Callable[..., tuple[dict[str, object], int]]
@ -22,10 +26,15 @@ def unwrap(func: object) -> ControllerMethod:
@pytest.fixture
def flask_app() -> Flask:
def flask_app() -> Iterator[Flask]:
app = Flask(__name__)
app.config["TESTING"] = True
return app
app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:"
db.init_app(app)
with app.app_context():
DataSourceOauthBinding.__table__.create(db.engine)
yield app
@pytest.fixture
@ -35,9 +44,13 @@ def current_user() -> Account:
return account
def test_get_data_source_integrates_serializes_orm_binding(flask_app: Flask) -> None:
TENANT_ID = "11111111-1111-1111-1111-111111111111"
BINDING_ID = "22222222-2222-2222-2222-222222222222"
def _add_binding(session: Session, *, disabled: bool) -> DataSourceOauthBinding:
binding = DataSourceOauthBinding(
tenant_id="tenant-1",
tenant_id=TENANT_ID,
access_token="token",
provider="notion",
source_info={
@ -55,24 +68,31 @@ def test_get_data_source_integrates_serializes_orm_binding(flask_app: Flask) ->
}
],
},
disabled=disabled,
)
binding.id = "binding-1"
binding.id = BINDING_ID
binding.created_at = datetime(2026, 5, 25, 1, 2, 3, tzinfo=UTC)
binding.disabled = False
session.add(binding)
session.commit()
return binding
with (
flask_app.test_request_context("/"),
patch.object(module.db.session, "scalars", return_value=MagicMock(all=lambda: [binding])),
):
response, status = unwrap(DataSourceApi().get)(DataSourceApi(), "tenant-1")
def test_get_data_source_integrates_serializes_orm_binding(
flask_app: Flask,
) -> None:
binding = _add_binding(db.session, disabled=False)
expected_created_at = int(binding.created_at.timestamp())
with flask_app.test_request_context("/"):
response, status = unwrap(DataSourceApi().get)(DataSourceApi(), TENANT_ID)
assert status == 200
assert response == {
"data": [
{
"id": "binding-1",
"id": BINDING_ID,
"provider": "notion",
"created_at": 1779670923,
"created_at": expected_created_at,
"is_bound": True,
"disabled": False,
"source_info": {
@ -96,34 +116,75 @@ def test_get_data_source_integrates_serializes_orm_binding(flask_app: Flask) ->
}
def test_get_data_source_integrates_preserves_empty_list_when_no_binding(flask_app: Flask) -> None:
with (
flask_app.test_request_context("/"),
patch.object(module.db.session, "scalars", return_value=MagicMock(all=lambda: [])),
):
response, status = unwrap(DataSourceApi().get)(DataSourceApi(), "tenant-1")
def test_get_data_source_integrates_preserves_empty_list_when_no_binding(
flask_app: Flask,
) -> None:
with flask_app.test_request_context("/"):
response, status = unwrap(DataSourceApi().get)(DataSourceApi(), TENANT_ID)
assert status == 200
assert response == {"data": []}
def test_patch_data_source_binding_uses_injected_session(flask_app: Flask) -> None:
binding = MagicMock(disabled=True)
session = MagicMock()
session.scalar.return_value = binding
@pytest.mark.parametrize(
("disabled", "action", "expected_disabled"),
[(True, "enable", False), (False, "disable", True)],
)
@pytest.mark.parametrize("sqlite_session", [(DataSourceOauthBinding,)], indirect=True)
def test_patch_data_source_binding_updates_state(
flask_app: Flask,
sqlite_session: Session,
disabled: bool,
action: Literal["enable", "disable"],
expected_disabled: bool,
) -> None:
_add_binding(sqlite_session, disabled=disabled)
sqlite_session.expunge_all()
with flask_app.test_request_context("/"):
response, status = unwrap(DataSourceApi().patch)(DataSourceApi(), session, "tenant-1", uuid4(), "enable")
response, status = unwrap(DataSourceApi().patch)(
DataSourceApi(), sqlite_session, TENANT_ID, UUID(BINDING_ID), action
)
sqlite_session.flush()
sqlite_session.expire_all()
binding = sqlite_session.scalar(select(DataSourceOauthBinding).where(DataSourceOauthBinding.id == BINDING_ID))
assert status == 200
assert response == {"result": "success"}
assert binding.disabled is False
session.scalar.assert_called_once()
session.add.assert_not_called()
session.commit.assert_not_called()
assert binding is not None
assert binding.disabled is expected_disabled
def test_notion_pre_import_pages_serializes_frontend_list_shape(flask_app: Flask, current_user: Account) -> None:
@pytest.mark.parametrize("sqlite_session", [(DataSourceOauthBinding,)], indirect=True)
def test_patch_data_source_binding_rejects_unknown_binding(
flask_app: Flask,
sqlite_session: Session,
) -> None:
with flask_app.test_request_context("/"), pytest.raises(NotFound, match="Data source binding not found"):
unwrap(DataSourceApi().patch)(DataSourceApi(), sqlite_session, TENANT_ID, UUID(BINDING_ID), "enable")
@pytest.mark.parametrize(("disabled", "action"), [(False, "enable"), (True, "disable")])
@pytest.mark.parametrize("sqlite_session", [(DataSourceOauthBinding,)], indirect=True)
def test_patch_data_source_binding_rejects_current_state(
flask_app: Flask,
sqlite_session: Session,
disabled: bool,
action: Literal["enable", "disable"],
) -> None:
_add_binding(sqlite_session, disabled=disabled)
sqlite_session.expunge_all()
with flask_app.test_request_context("/"), pytest.raises(ValueError):
unwrap(DataSourceApi().patch)(DataSourceApi(), sqlite_session, TENANT_ID, UUID(BINDING_ID), action)
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
def test_notion_pre_import_pages_serializes_frontend_list_shape(
flask_app: Flask,
current_user: Account,
sqlite_session: Session,
) -> None:
page = MagicMock(
page_id="page-1",
page_name="Page",
@ -145,8 +206,6 @@ def test_notion_pre_import_pages_serializes_frontend_list_shape(flask_app: Flask
get_online_document_pages=MagicMock(return_value=iter([online_document_message])),
datasource_provider_type=MagicMock(return_value="online_document"),
)
session = MagicMock()
with (
flask_app.test_request_context("/?credential_id=credential-1"),
patch.object(
@ -158,7 +217,7 @@ def test_notion_pre_import_pages_serializes_frontend_list_shape(flask_app: Flask
patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime),
):
response, status = unwrap(DataSourceNotionListApi().get)(
DataSourceNotionListApi(), session, "tenant-1", current_user
DataSourceNotionListApi(), sqlite_session, "tenant-1", current_user
)
assert status == 200
@ -183,3 +242,38 @@ def test_notion_pre_import_pages_serializes_frontend_list_shape(flask_app: Flask
}
runtime.get_online_document_pages.assert_called_once()
assert runtime.get_online_document_pages.call_args.kwargs["datasource_parameters"] == {}
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
def test_notion_pre_import_pages_rejects_missing_credential(
flask_app: Flask,
current_user: Account,
sqlite_session: Session,
) -> None:
with (
flask_app.test_request_context("/?credential_id=credential-1"),
patch.object(module.DatasourceProviderService, "get_datasource_credentials", return_value=None),
pytest.raises(NotFound, match="Credential not found"),
):
unwrap(DataSourceNotionListApi().get)(DataSourceNotionListApi(), sqlite_session, TENANT_ID, current_user)
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
def test_notion_pre_import_pages_rejects_non_notion_dataset(
flask_app: Flask,
current_user: Account,
sqlite_session: Session,
) -> None:
dataset = MagicMock(data_source_type="other_type")
with (
flask_app.test_request_context("/?credential_id=credential-1&dataset_id=dataset-1"),
patch.object(
module.DatasourceProviderService,
"get_datasource_credentials",
return_value={"token": "token"},
),
patch.object(module.DatasetService, "get_dataset", return_value=dataset),
pytest.raises(ValueError, match="Dataset is not notion type"),
):
unwrap(DataSourceNotionListApi().get)(DataSourceNotionListApi(), sqlite_session, TENANT_ID, current_user)

171
api/tests/unit_tests/controllers/console/datasets/test_data_source_notion_apis.py Normal file
View File

@ -0,0 +1,171 @@
"""Unit tests for controllers.console.datasets.data_source Notion endpoints."""
from __future__ import annotations
import inspect
from unittest.mock import MagicMock, patch
import pytest
from flask import Flask
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound
from controllers.console.datasets.data_source import (
DataSourceNotionDatasetSyncApi,
DataSourceNotionDocumentSyncApi,
DataSourceNotionIndexingEstimateApi,
DataSourceNotionPreviewApi,
)
from core.rag.index_processor.constant.index_type import IndexStructureType
from models import Account
@pytest.fixture
def current_user() -> Account:
account = Account(name="Test User", email="u1@example.com")
account.id = "u1"
return account
class TestDataSourceNotionPreviewApi:
def test_get_preview_success(self, app: Flask) -> None:
api = DataSourceNotionPreviewApi()
method = inspect.unwrap(api.get)
extractor = MagicMock(extract=lambda: [MagicMock(page_content="hello")])
with (
app.test_request_context("/?credential_id=c1"),
patch(
"controllers.console.datasets.data_source.DatasourceProviderService.get_datasource_credentials",
return_value={"integration_secret": "t"},
),
patch(
"controllers.console.datasets.data_source.NotionExtractor",
return_value=extractor,
),
):
response, status = method(api, "tenant-1", "p1", "page")
assert status == 200
class TestDataSourceNotionIndexingEstimateApi:
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
def test_post_indexing_estimate_success(self, app: Flask, sqlite_session: Session) -> None:
api = DataSourceNotionIndexingEstimateApi()
method = inspect.unwrap(api.post)
empty_rules: dict[str, object] = {}
payload: dict[str, object] = {
"notion_info_list": [
{
"workspace_id": "w1",
"credential_id": "c1",
"pages": [{"page_id": "p1", "type": "page"}],
}
],
"process_rule": {"rules": empty_rules},
"doc_form": IndexStructureType.PARAGRAPH_INDEX,
"doc_language": "English",
}
with (
app.test_request_context("/", method="POST", json=payload, headers={"Content-Type": "application/json"}),
patch(
"controllers.console.datasets.data_source.DocumentService.estimate_args_validate",
),
patch(
"controllers.console.datasets.data_source.IndexingRunner.indexing_estimate",
return_value=MagicMock(model_dump=lambda: {"total_pages": 1}),
),
):
response, status = method(api, sqlite_session, "tenant-1")
assert status == 200
class TestDataSourceNotionDatasetSyncApi:
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
def test_get_success(self, app: Flask, sqlite_session: Session) -> None:
api = DataSourceNotionDatasetSyncApi()
method = inspect.unwrap(api.get)
with (
app.test_request_context("/"),
patch(
"controllers.console.datasets.data_source.DatasetService.get_dataset",
return_value=MagicMock(),
),
patch(
"controllers.console.datasets.data_source.DocumentService.get_document_by_dataset_id",
return_value=[MagicMock(id="d1")],
),
patch(
"controllers.console.datasets.data_source.document_indexing_sync_task.delay",
return_value=None,
),
):
response, status = method(api, sqlite_session, "ds-1")
assert status == 200
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
def test_get_dataset_not_found(self, app: Flask, sqlite_session: Session) -> None:
api = DataSourceNotionDatasetSyncApi()
method = inspect.unwrap(api.get)
with (
app.test_request_context("/"),
patch(
"controllers.console.datasets.data_source.DatasetService.get_dataset",
return_value=None,
),
):
with pytest.raises(NotFound):
method(api, sqlite_session, "ds-1")
class TestDataSourceNotionDocumentSyncApi:
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
def test_get_success(self, app: Flask, sqlite_session: Session) -> None:
api = DataSourceNotionDocumentSyncApi()
method = inspect.unwrap(api.get)
with (
app.test_request_context("/"),
patch(
"controllers.console.datasets.data_source.DatasetService.get_dataset",
return_value=MagicMock(),
),
patch(
"controllers.console.datasets.data_source.DocumentService.get_document",
return_value=MagicMock(),
),
patch(
"controllers.console.datasets.data_source.document_indexing_sync_task.delay",
return_value=None,
),
):
response, status = method(api, sqlite_session, "ds-1", "doc-1")
assert status == 200
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
def test_get_document_not_found(self, app: Flask, sqlite_session: Session) -> None:
api = DataSourceNotionDocumentSyncApi()
method = inspect.unwrap(api.get)
with (
app.test_request_context("/"),
patch(
"controllers.console.datasets.data_source.DatasetService.get_dataset",
return_value=MagicMock(),
),
patch(
"controllers.console.datasets.data_source.DocumentService.get_document",
return_value=None,
),
):
with pytest.raises(NotFound):
method(api, sqlite_session, "ds-1", "doc-1")