From e1ce808567574e755d141af5bf283b9e05601a1e Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 28 Jul 2026 12:51:22 +0900 Subject: [PATCH] test: move data source controller coverage to unit tests (#38924) --- .../console/datasets/test_data_source.py | 551 +++--------------- .../console/datasets/test_data_source.py | 164 ++++-- .../datasets/test_data_source_notion_apis.py | 171 ++++++ 3 files changed, 371 insertions(+), 515 deletions(-) create mode 100644 api/tests/unit_tests/controllers/console/datasets/test_data_source_notion_apis.py diff --git a/api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py b/api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py index 7c63cee4107..814886b1772 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py +++ b/api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py @@ -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 diff --git a/api/tests/unit_tests/controllers/console/datasets/test_data_source.py b/api/tests/unit_tests/controllers/console/datasets/test_data_source.py index a6cb79417a7..55ee174355c 100644 --- a/api/tests/unit_tests/controllers/console/datasets/test_data_source.py +++ b/api/tests/unit_tests/controllers/console/datasets/test_data_source.py @@ -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) diff --git a/api/tests/unit_tests/controllers/console/datasets/test_data_source_notion_apis.py b/api/tests/unit_tests/controllers/console/datasets/test_data_source_notion_apis.py new file mode 100644 index 00000000000..d86a1093cc8 --- /dev/null +++ b/api/tests/unit_tests/controllers/console/datasets/test_data_source_notion_apis.py @@ -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")