opensource/dify
3
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-21 13:20:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: replace manual model_validate with @model_validate in datasets controllers (#40236)

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Likalikali
2026-08-09 08:25:27 +00:00
committed by GitHub
co-authored by autofix-ci[bot]
parent b692ddd80c
commit e09eee8fd1
12 changed files with 406 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py
+6 -2
View File
@@ -7,7 +7,7 @@ from uuid import uuid4
from flask import Flask
from sqlalchemy.orm import Session
from controllers.console.datasets.data_source import DataSourceNotionListApi
from controllers.console.datasets.data_source import DataSourceNotionListApi, DataSourceNotionListQuery
from models import Account
from models.dataset import Document
from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus
@@ -78,7 +78,11 @@ def test_notion_page_is_marked_bound_from_persisted_document(
),
):
response, status = unwrap(DataSourceNotionListApi().get)(
DataSourceNotionListApi(), db_session_with_containers, tenant_id, account
DataSourceNotionListApi(),
DataSourceNotionListQuery(credential_id="c1", dataset_id=dataset_id),
db_session_with_containers,
tenant_id,
account,
)
assert status == 200
api/tests/unit_tests/controllers/console/datasets/test_data_source.py
+20 -4
View File
@@ -14,7 +14,7 @@ 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 controllers.console.datasets.data_source import DataSourceApi, DataSourceNotionListApi, DataSourceNotionListQuery
from models import Account, DataSourceOauthBinding
from models.engine import db
@@ -217,7 +217,11 @@ def test_notion_pre_import_pages_serializes_frontend_list_shape(
patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime),
):
response, status = unwrap(DataSourceNotionListApi().get)(
DataSourceNotionListApi(), sqlite_session, "tenant-1", current_user
DataSourceNotionListApi(),
DataSourceNotionListQuery(credential_id="credential-1"),
sqlite_session,
"tenant-1",
current_user,
)
assert status == 200
@@ -255,7 +259,13 @@ def test_notion_pre_import_pages_rejects_missing_credential(
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)
unwrap(DataSourceNotionListApi().get)(
DataSourceNotionListApi(),
DataSourceNotionListQuery(credential_id="credential-1"),
sqlite_session,
TENANT_ID,
current_user,
)
@pytest.mark.parametrize("sqlite_session", [()], indirect=True)
@@ -276,4 +286,10 @@ def test_notion_pre_import_pages_rejects_non_notion_dataset(
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)
unwrap(DataSourceNotionListApi().get)(
DataSourceNotionListApi(),
DataSourceNotionListQuery(credential_id="credential-1", dataset_id="dataset-1"),
sqlite_session,
TENANT_ID,
current_user,
)
api/tests/unit_tests/controllers/console/datasets/test_data_source_notion_apis.py
+4 -2
View File
@@ -15,6 +15,8 @@ from controllers.console.datasets.data_source import (
DataSourceNotionDocumentSyncApi,
DataSourceNotionIndexingEstimateApi,
DataSourceNotionPreviewApi,
DataSourceNotionPreviewQuery,
NotionEstimatePayload,
)
from core.rag.index_processor.constant.index_type import IndexStructureType
from models import Account
@@ -45,7 +47,7 @@ class TestDataSourceNotionPreviewApi:
return_value=extractor,
),
):
response, status = method(api, "tenant-1", "p1", "page")
response, status = method(api, DataSourceNotionPreviewQuery(credential_id="c1"), "tenant-1", "p1", "page")
assert status == 200
@@ -80,7 +82,7 @@ class TestDataSourceNotionIndexingEstimateApi:
return_value=MagicMock(model_dump=lambda: {"total_pages": 1}),
),
):
response, status = method(api, sqlite_session, "tenant-1")
response, status = method(api, NotionEstimatePayload.model_validate(payload), sqlite_session, "tenant-1")
assert status == 200
api/tests/unit_tests/controllers/console/datasets/test_datasets.py
+44 -16
View File
@@ -18,6 +18,7 @@ from controllers.console.datasets.datasets import (
DatasetApiDeleteApi,
DatasetApiKeyApi,
DatasetAutoDisableLogApi,
DatasetCreatePayload,
DatasetEnableApiApi,
DatasetErrorDocs,
DatasetIndexingEstimateApi,
@@ -28,7 +29,9 @@ from controllers.console.datasets.datasets import (
DatasetRelatedAppListApi,
DatasetRetrievalSettingApi,
DatasetRetrievalSettingMockApi,
DatasetUpdatePayload,
DatasetUseCheckApi,
IndexingEstimatePayload,
_get_retrieval_methods_by_vector_type,
)
from controllers.console.datasets.error import DatasetInUseError, DatasetNameDuplicateError, IndexingEstimateError
@@ -489,7 +492,7 @@ class TestDatasetListApiPost:
patch.object(type(console_ns), "payload", payload),
patch.object(DatasetService, "create_empty_dataset", return_value=dataset),
):
_, status = method(api, MagicMock(), "tenant-1", user)
_, status = method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", user)
assert status == 201
def test_post_forbidden(self, app: Flask):
@@ -499,7 +502,7 @@ class TestDatasetListApiPost:
user = make_account(TenantAccountRole.NORMAL)
with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload):
with pytest.raises(Forbidden):
method(api, MagicMock(), "tenant-1", user)
method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", user)
def test_post_duplicate_name(self, app: Flask):
api = DatasetListApi()
@@ -514,14 +517,14 @@ class TestDatasetListApiPost:
),
):
with pytest.raises(DatasetNameDuplicateError):
method(api, MagicMock(), "tenant-1", user)
method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", user)
def test_post_invalid_payload_missing_name(self, app: Flask):
api = DatasetListApi()
method = unwrap(api.post)
with app.test_request_context("/datasets", json={}), patch.object(type(console_ns), "payload", {}):
with pytest.raises(ValueError):
method(api, MagicMock(), "tenant-1", make_account())
method(api, DatasetCreatePayload(), MagicMock(), "tenant-1", make_account())
def test_post_invalid_indexing_technique(self, app: Flask):
api = DatasetListApi()
@@ -529,7 +532,7 @@ class TestDatasetListApiPost:
payload = {"name": "bad", "indexing_technique": "invalid-tech"}
with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload):
with pytest.raises(ValueError, match="Invalid indexing technique"):
method(api, MagicMock(), "tenant-1", make_account())
method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", make_account())
def test_post_invalid_provider(self, app: Flask):
api = DatasetListApi()
@@ -537,7 +540,7 @@ class TestDatasetListApiPost:
payload = {"name": "bad", "provider": "unknown"}
with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload):
with pytest.raises(ValueError, match="Invalid provider"):
method(api, MagicMock(), "tenant-1", make_account())
method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", make_account())
class TestDatasetApiGet:
@@ -692,7 +695,7 @@ class TestDatasetApiPatch:
patch.object(DatasetService, "update_dataset", return_value=dataset),
patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=[]),
):
result, status = method(api, MagicMock(), tenant_id, user, dataset_id)
result, status = method(api, DatasetUpdatePayload(), MagicMock(), tenant_id, user, dataset_id)
assert status == 200
assert result["partial_member_list"] == []
@@ -704,7 +707,7 @@ class TestDatasetApiPatch:
patch.object(DatasetService, "get_dataset", return_value=None),
):
with pytest.raises(NotFound, match="Dataset not found"):
method(api, MagicMock(), "tenant-1", make_account(), "missing")
method(api, DatasetUpdatePayload(), MagicMock(), "tenant-1", make_account(), "missing")
def test_patch_permission_denied(self, app: Flask):
api = DatasetApi()
@@ -719,7 +722,7 @@ class TestDatasetApiPatch:
patch.object(DatasetPermissionService, "check_permission", side_effect=Forbidden("no permission")),
):
with pytest.raises(Forbidden):
method(api, MagicMock(), "tenant", make_account(), dataset_id)
method(api, DatasetUpdatePayload(), MagicMock(), "tenant", make_account(), dataset_id)
def test_patch_partial_members_update(self, app: Flask):
api = DatasetApi()
@@ -736,7 +739,7 @@ class TestDatasetApiPatch:
patch.object(DatasetPermissionService, "update_partial_member_list", return_value=None),
patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=["u1", "u2"]),
):
result, _ = method(api, MagicMock(), "tenant", make_account(), dataset_id)
result, _ = method(api, DatasetUpdatePayload(), MagicMock(), "tenant", make_account(), dataset_id)
assert result["partial_member_list"] == ["u1", "u2"]
def test_patch_clear_partial_members(self, app: Flask):
@@ -754,7 +757,7 @@ class TestDatasetApiPatch:
patch.object(DatasetPermissionService, "clear_partial_member_list", return_value=None),
patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=[]),
):
result, _ = method(api, MagicMock(), "tenant", make_account(), dataset_id)
result, _ = method(api, DatasetUpdatePayload(), MagicMock(), "tenant", make_account(), dataset_id)
assert result["partial_member_list"] == []
@@ -1014,7 +1017,12 @@ class TestDatasetIndexingEstimateApi:
patch("controllers.console.datasets.datasets.DocumentService.estimate_args_validate", return_value=None),
patch("controllers.console.datasets.datasets.IndexingRunner.indexing_estimate", return_value=mock_response),
):
response, status = method(api, session, "tenant-1")
response, status = method(
api,
IndexingEstimatePayload(**payload),
session,
"tenant-1",
)
assert status == 200
assert response == {
"tokens": 0,
@@ -1036,7 +1044,12 @@ class TestDatasetIndexingEstimateApi:
patch("controllers.console.datasets.datasets.DocumentService.estimate_args_validate", return_value=None),
):
with pytest.raises(NotFound):
method(api, session, "tenant-1")
method(
api,
IndexingEstimatePayload(**payload),
session,
"tenant-1",
)
def test_post_llm_bad_request_error(self, app: Flask):
api = DatasetIndexingEstimateApi()
@@ -1055,7 +1068,12 @@ class TestDatasetIndexingEstimateApi:
),
):
with pytest.raises(ProviderNotInitializeError):
method(api, session, "tenant-1")
method(
api,
IndexingEstimatePayload(**payload),
session,
"tenant-1",
)
def test_post_provider_token_not_init(self, app: Flask):
api = DatasetIndexingEstimateApi()
@@ -1074,7 +1092,12 @@ class TestDatasetIndexingEstimateApi:
),
):
with pytest.raises(ProviderNotInitializeError):
method(api, session, "tenant-1")
method(
api,
IndexingEstimatePayload(**payload),
session,
"tenant-1",
)
def test_post_generic_exception(self, app: Flask):
api = DatasetIndexingEstimateApi()
@@ -1092,7 +1115,12 @@ class TestDatasetIndexingEstimateApi:
),
):
with pytest.raises(IndexingEstimateError):
method(api, session, "tenant-1")
method(
api,
IndexingEstimatePayload(**payload),
session,
"tenant-1",
)
class TestDatasetRelatedAppListApi:
api/tests/unit_tests/controllers/console/datasets/test_datasets_segments.py
+129 -21
View File
@@ -9,9 +9,11 @@ from flask import Flask
from werkzeug.exceptions import Forbidden, NotFound
import services
from controllers.common.controller_schemas import ChildChunkCreatePayload, ChildChunkUpdatePayload
from controllers.console import console_ns
from controllers.console.app.error import ProviderNotInitializeError
from controllers.console.datasets.datasets_segments import (
BatchImportPayload,
ChildChunkAddApi,
ChildChunkBatchUpdatePayload,
ChildChunkUpdateApi,
@@ -20,6 +22,8 @@ from controllers.console.datasets.datasets_segments import (
DatasetDocumentSegmentBatchImportApi,
DatasetDocumentSegmentListApi,
DatasetDocumentSegmentUpdateApi,
SegmentCreatePayload,
SegmentUpdatePayload,
)
from controllers.console.datasets.error import ChildChunkDeleteIndexError, ChildChunkIndexingError, InvalidActionError
from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
@@ -50,6 +54,7 @@ def _segment():
status=SegmentStatus.COMPLETED,
updated_by="u1",
)
segment.id = "seg-1"
segment.created_at = naive_utc_now()
segment.updated_at = naive_utc_now()
@@ -354,7 +359,9 @@ class TestDatasetDocumentSegmentAddApi:
return_value=None,
),
):
response, status = method(api, session, "tenant-1", user, "ds-1", "doc-1")
response, status = method(
api, SegmentCreatePayload(content="test content"), session, "tenant-1", user, "ds-1", "doc-1"
)
assert status == 200
assert response["data"]["id"] == "seg-1"
@@ -378,7 +385,9 @@ class TestDatasetDocumentSegmentAddApi:
),
):
with pytest.raises(ProviderNotInitializeError):
method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1")
method(
api, SegmentCreatePayload(content="test content"), MagicMock(), "tenant-1", user, "ds-1", "doc-1"
)
def test_post_provider_token_not_init(self, app: Flask):
api = DatasetDocumentSegmentAddApi()
@@ -400,7 +409,9 @@ class TestDatasetDocumentSegmentAddApi:
),
):
with pytest.raises(ProviderNotInitializeError):
method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1")
method(
api, SegmentCreatePayload(content="test content"), MagicMock(), "tenant-1", user, "ds-1", "doc-1"
)
class TestDatasetDocumentSegmentUpdateApi:
@@ -441,7 +452,9 @@ class TestDatasetDocumentSegmentUpdateApi:
return_value=None,
),
):
response, status = method(api, session, "tenant-1", user, "ds-1", "doc-1", "seg-1")
response, status = method(
api, SegmentUpdatePayload(content="test content"), session, "tenant-1", user, "ds-1", "doc-1", "seg-1"
)
assert status == 200
assert "data" in response
@@ -467,7 +480,16 @@ class TestDatasetDocumentSegmentUpdateApi:
),
):
with pytest.raises(NotFound):
method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1")
method(
api,
SegmentUpdatePayload(content="test content"),
MagicMock(),
"tenant-1",
user,
"ds-1",
"doc-1",
"seg-1",
)
def test_patch_segment_not_found(self, app: Flask):
api = DatasetDocumentSegmentUpdateApi()
@@ -495,7 +517,16 @@ class TestDatasetDocumentSegmentUpdateApi:
),
):
with pytest.raises(NotFound):
method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1")
method(
api,
SegmentUpdatePayload(content="test content"),
MagicMock(),
"tenant-1",
user,
"ds-1",
"doc-1",
"seg-1",
)
def test_patch_llm_bad_request(self, app: Flask):
api = DatasetDocumentSegmentUpdateApi()
@@ -525,7 +556,16 @@ class TestDatasetDocumentSegmentUpdateApi:
),
):
with pytest.raises(ProviderNotInitializeError):
method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1")
method(
api,
SegmentUpdatePayload(content="test content"),
MagicMock(),
"tenant-1",
user,
"ds-1",
"doc-1",
"seg-1",
)
class TestDatasetDocumentSegmentBatchImportApi:
@@ -564,7 +604,9 @@ class TestDatasetDocumentSegmentBatchImportApi:
return_value=None,
),
):
response, status = method(api, session, "tenant-1", user, "ds-1", "doc-1")
response, status = method(
api, BatchImportPayload(upload_file_id="test-file-id"), session, "tenant-1", user, "ds-1", "doc-1"
)
assert status == 200
assert response["job_status"] == "waiting"
@@ -581,7 +623,15 @@ class TestDatasetDocumentSegmentBatchImportApi:
patch("controllers.console.datasets.datasets_segments.DatasetService.get_dataset", return_value=None),
):
with pytest.raises(NotFound):
method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1")
method(
api,
BatchImportPayload(upload_file_id="test-file-id"),
MagicMock(),
"tenant-1",
user,
"ds-1",
"doc-1",
)
def test_post_document_not_found(self, app: Flask):
api = DatasetDocumentSegmentBatchImportApi()
@@ -599,7 +649,15 @@ class TestDatasetDocumentSegmentBatchImportApi:
patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=None),
):
with pytest.raises(NotFound):
method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1")
method(
api,
BatchImportPayload(upload_file_id="test-file-id"),
MagicMock(),
"tenant-1",
user,
"ds-1",
"doc-1",
)
def test_post_upload_file_not_found(self, app: Flask):
api = DatasetDocumentSegmentBatchImportApi()
@@ -619,7 +677,9 @@ class TestDatasetDocumentSegmentBatchImportApi:
),
):
with pytest.raises(NotFound):
method(api, session, "tenant-1", user, "ds-1", "doc-1")
method(
api, BatchImportPayload(upload_file_id="test-file-id"), session, "tenant-1", user, "ds-1", "doc-1"
)
def test_post_invalid_file_type(self, app: Flask):
api = DatasetDocumentSegmentBatchImportApi()
@@ -641,7 +701,9 @@ class TestDatasetDocumentSegmentBatchImportApi:
),
):
with pytest.raises(ValueError):
method(api, session, "tenant-1", user, "ds-1", "doc-1")
method(
api, BatchImportPayload(upload_file_id="test-file-id"), session, "tenant-1", user, "ds-1", "doc-1"
)
def test_post_async_task_failure(self, app: Flask):
api = DatasetDocumentSegmentBatchImportApi()
@@ -665,7 +727,9 @@ class TestDatasetDocumentSegmentBatchImportApi:
"controllers.console.datasets.datasets_segments.redis_client.setnx", side_effect=Exception("redis down")
),
):
response, status = method(api, session, "tenant-1", user, "ds-1", "doc-1")
response, status = method(
api, BatchImportPayload(upload_file_id="test-file-id"), session, "tenant-1", user, "ds-1", "doc-1"
)
assert status == 500
assert "error" in response
@@ -747,7 +811,9 @@ class TestChildChunkAddApi:
return_value=child_chunk,
),
):
response, status = method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1")
response, status = method(
api, ChildChunkCreatePayload(content="child"), MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1"
)
assert status == 200
assert response["data"]["id"] == "cc-1"
@@ -778,7 +844,16 @@ class TestChildChunkAddApi:
),
):
with pytest.raises(ChildChunkIndexingError):
method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1")
method(
api,
ChildChunkCreatePayload(content="child"),
MagicMock(),
"tenant-1",
user,
"ds-1",
"doc-1",
"seg-1",
)
def test_post_permission_denied(self, app: Flask):
api = ChildChunkAddApi()
@@ -799,7 +874,16 @@ class TestChildChunkAddApi:
),
):
with pytest.raises(Forbidden):
method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1")
method(
api,
ChildChunkCreatePayload(content="child"),
MagicMock(),
"tenant-1",
user,
"ds-1",
"doc-1",
"seg-1",
)
class TestChildChunkUpdateApi:
@@ -922,7 +1006,17 @@ class TestChildChunkUpdateApi:
),
):
with pytest.raises(NotFound):
method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1", "seg-1", "cc-1")
method(
api,
ChildChunkUpdatePayload(content="updated child"),
MagicMock(),
"tenant-1",
user,
"ds-1",
"doc-1",
"seg-1",
"cc-1",
)
class TestSegmentListAdvancedCases:
@@ -1047,7 +1141,9 @@ class TestSegmentOperationCases:
),
):
with pytest.raises(ProviderTokenNotInitError):
method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1")
method(
api, SegmentCreatePayload(content="test content"), MagicMock(), "tenant-1", user, "ds-1", "doc-1"
)
def test_batch_import_with_document_not_found(self, app: Flask):
"""Test batch import with document not found"""
@@ -1063,7 +1159,15 @@ class TestSegmentOperationCases:
patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=None),
):
with pytest.raises(NotFound):
method(api, MagicMock(), "tenant-1", user, "ds-1", "doc-1")
method(
api,
BatchImportPayload(upload_file_id="test-file-id"),
MagicMock(),
"tenant-1",
user,
"ds-1",
"doc-1",
)
def test_batch_import_with_invalid_file(self, app: Flask):
"""Test batch import with invalid file type"""
@@ -1083,7 +1187,9 @@ class TestSegmentOperationCases:
patch("controllers.console.datasets.datasets_segments.DocumentService.get_document", return_value=document),
):
with pytest.raises(NotFound):
method(api, session, "tenant-1", user, "ds-1", "doc-1")
method(
api, BatchImportPayload(upload_file_id="test-file-id"), session, "tenant-1", user, "ds-1", "doc-1"
)
def test_batch_import_with_async_task_failure(self, app: Flask):
api = DatasetDocumentSegmentBatchImportApi()
@@ -1122,7 +1228,9 @@ class TestSegmentOperationCases:
side_effect=Exception("Task failed"),
),
):
response, status = method(api, session, "tenant-1", user, "ds-1", "doc-1")
response, status = method(
api, BatchImportPayload(upload_file_id="test-file-id"), session, "tenant-1", user, "ds-1", "doc-1"
)
assert status == 500
assert "error" in response
api/tests/unit_tests/controllers/console/datasets/test_external.py
+33 -16
View File
@@ -12,10 +12,14 @@ from controllers.console import console_ns
from controllers.console.datasets.error import DatasetNameDuplicateError
from controllers.console.datasets.external import (
BedrockRetrievalApi,
BedrockRetrievalPayload,
ExternalApiTemplateApi,
ExternalApiTemplateListApi,
ExternalApiTemplateListQuery,
ExternalApiUseCheckApi,
ExternalDatasetCreateApi,
ExternalHitTestingPayload,
ExternalKnowledgeApiPayload,
ExternalKnowledgeHitTestingApi,
)
from models.account import Account, TenantAccountRole
@@ -174,7 +178,9 @@ class TestExternalApiTemplateListApi:
return_value=([api_item], 3),
) as get_external_knowledge_apis,
):
resp, status = method(api, session, "tenant-1")
resp, status = method(
api, ExternalApiTemplateListQuery(page=2, limit=1, keyword="vector"), session, "tenant-1"
)
assert status == 200
assert resp == {
@@ -213,7 +219,9 @@ class TestExternalApiTemplateListApi:
return_value=created,
) as create_external_knowledge_api,
):
resp, status = method(api, session, "tenant-1", current_user)
resp, status = method(
api, ExternalKnowledgeApiPayload.model_validate(payload), session, "tenant-1", current_user
)
assert status == 201
assert resp == _external_api_dict("api-created")
@@ -239,7 +247,7 @@ class TestExternalApiTemplateListApi:
patch.object(ExternalDatasetService, "validate_api_list"),
):
with pytest.raises(Forbidden):
method(api, MagicMock(), "tenant-1", current_user)
method(api, ExternalKnowledgeApiPayload.model_validate(payload), MagicMock(), "tenant-1", current_user)
def test_post_duplicate_name(self, app: Flask, current_user: Account):
api = ExternalApiTemplateListApi()
@@ -258,7 +266,7 @@ class TestExternalApiTemplateListApi:
),
):
with pytest.raises(DatasetNameDuplicateError):
method(api, MagicMock(), "tenant-1", current_user)
method(api, ExternalKnowledgeApiPayload.model_validate(payload), MagicMock(), "tenant-1", current_user)
class TestExternalApiTemplateApi:
@@ -325,7 +333,14 @@ class TestExternalApiTemplateApi:
return_value=updated,
) as update_external_knowledge_api,
):
resp, status = method(api, session, "tenant-1", current_user, "api-updated")
resp, status = method(
api,
ExternalKnowledgeApiPayload.model_validate(payload),
session,
"tenant-1",
current_user,
"api-updated",
)
assert status == 200
assert resp == _external_api_dict("api-updated")
@@ -405,7 +420,9 @@ class TestExternalDatasetCreateApi:
) as dataset_response_source,
):
session = MagicMock()
resp, status = method(api, session, "tenant-1", current_user)
resp, status = method(
api, ExternalDatasetCreatePayload.model_validate(payload), session, "tenant-1", current_user
)
assert status == 201
assert resp == _expected_dataset_detail_payload()
@@ -433,7 +450,7 @@ class TestExternalDatasetCreateApi:
patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload),
):
with pytest.raises(Forbidden):
method(api, MagicMock(), "tenant-1", current_user)
method(api, ExternalDatasetCreatePayload.model_validate(payload), MagicMock(), "tenant-1", current_user)
class TestExternalKnowledgeHitTestingApi:
@@ -450,7 +467,7 @@ class TestExternalKnowledgeHitTestingApi:
),
):
with pytest.raises(NotFound):
method(api, MagicMock(), current_user, "dataset-id")
method(api, ExternalHitTestingPayload(query="test"), MagicMock(), current_user, "dataset-id")
def test_hit_testing_success(self, app: Flask, current_user: Account):
api = ExternalKnowledgeHitTestingApi()
@@ -496,7 +513,7 @@ class TestExternalKnowledgeHitTestingApi:
) as external_retrieve,
patch("controllers.console.datasets.external.dump_response", side_effect=lambda _model, value: value),
):
resp = method(api, session, current_user, "dataset-id")
resp = method(api, ExternalHitTestingPayload.model_validate(payload), session, current_user, "dataset-id")
assert resp == retrieve_response
check_dataset_permission.assert_called_once_with(dataset, current_user, session)
@@ -549,7 +566,7 @@ class TestBedrockRetrievalApi:
return_value=retrieval_response,
) as knowledge_retrieval,
):
resp, status = method()
resp, status = method(api, BedrockRetrievalPayload.model_validate(payload))
assert status == 200
assert resp == retrieval_response
@@ -576,7 +593,7 @@ class TestExternalApiTemplateListApiAdvanced:
),
):
with pytest.raises(DatasetNameDuplicateError):
method(api, MagicMock(), "tenant-1", current_user)
method(api, ExternalKnowledgeApiPayload.model_validate(payload), MagicMock(), "tenant-1", current_user)
def test_get_with_pagination(self, app: Flask):
api = ExternalApiTemplateListApi()
@@ -591,7 +608,7 @@ class TestExternalApiTemplateListApiAdvanced:
return_value=(templates, 25),
) as get_external_knowledge_apis,
):
resp, status = method(api, MagicMock(), "tenant-1")
resp, status = method(api, ExternalApiTemplateListQuery(page=2, limit=3), MagicMock(), "tenant-1")
assert status == 200
assert resp == {
@@ -621,7 +638,7 @@ class TestExternalDatasetCreateApiAdvanced:
with app.test_request_context("/", json=payload), patch.object(type(console_ns), "payload", payload):
with pytest.raises(Forbidden):
method(api, MagicMock(), "tenant-1", current_user)
method(api, ExternalDatasetCreatePayload.model_validate(payload), MagicMock(), "tenant-1", current_user)
class TestExternalKnowledgeHitTestingApiAdvanced:
@@ -644,7 +661,7 @@ class TestExternalKnowledgeHitTestingApiAdvanced:
),
):
with pytest.raises(NotFound):
method(api, MagicMock(), current_user, "ds-1")
method(api, ExternalHitTestingPayload.model_validate(payload), MagicMock(), current_user, "ds-1")
def test_hit_testing_with_custom_retrieval_model(self, app: Flask, current_user: Account):
api = ExternalKnowledgeHitTestingApi()
@@ -682,7 +699,7 @@ class TestExternalKnowledgeHitTestingApiAdvanced:
},
) as external_retrieve,
):
resp = method(api, session, current_user, "ds-1")
resp = method(api, ExternalHitTestingPayload.model_validate(payload), session, current_user, "ds-1")
assert resp == {
"query": {"content": "test query"},
@@ -727,4 +744,4 @@ class TestBedrockRetrievalApiAdvanced:
),
):
with pytest.raises(ValueError):
method()
method(api, BedrockRetrievalPayload.model_validate(payload))
api/tests/unit_tests/controllers/console/datasets/test_metadata.py
+24 -5
View File
@@ -7,6 +7,7 @@ from flask import Flask
from pytest_mock import MockerFixture
from werkzeug.exceptions import NotFound
from controllers.common.controller_schemas import MetadataUpdatePayload
from controllers.console import console_ns
from controllers.console.datasets.metadata import (
DatasetMetadataApi,
@@ -24,6 +25,7 @@ from services.metadata_service import MetadataService
@pytest.fixture
def app():
app = Flask("test_dataset_metadata")
app.config["TESTING"] = True
return app
@@ -76,7 +78,9 @@ class TestDatasetMetadataCreateApi:
MetadataService, "create_metadata", return_value={"id": "m1", "type": "string", "name": "author"}
),
):
result, status = method(api, MagicMock(), "tenant-1", current_user, dataset_id)
result, status = method(
api, MetadataArgs(type="string", name="author"), MagicMock(), "tenant-1", current_user, dataset_id
)
assert status == 201
assert result["type"] == "string"
assert result["name"] == "author"
@@ -92,7 +96,9 @@ class TestDatasetMetadataCreateApi:
patch.object(DatasetService, "get_dataset", return_value=None),
):
with pytest.raises(NotFound, match="Dataset not found"):
method(api, MagicMock(), "tenant-1", current_user, dataset_id)
method(
api, MetadataArgs(type="string", name="author"), MagicMock(), "tenant-1", current_user, dataset_id
)
class TestDatasetMetadataGetApi:
@@ -140,7 +146,15 @@ class TestDatasetMetadataApi:
return_value={"id": "m1", "type": "string", "name": "updated-name"},
),
):
result, status = method(api, MagicMock(), "tenant-1", current_user, dataset_id, metadata_id)
result, status = method(
api,
MetadataUpdatePayload(name="updated-name"),
MagicMock(),
"tenant-1",
current_user,
dataset_id,
metadata_id,
)
assert status == 200
assert result["type"] == "string"
assert result["name"] == "updated-name"
@@ -201,9 +215,14 @@ class TestDocumentMetadataEditApi:
patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload),
patch.object(DatasetService, "get_dataset", return_value=dataset),
patch.object(DatasetService, "check_dataset_permission"),
patch.object(MetadataOperationData, "model_validate", return_value=MagicMock()),
patch.object(MetadataService, "update_documents_metadata"),
):
result, status = method(api, MagicMock(), current_user, dataset_id)
result, status = method(
api,
MetadataOperationData(operation_data=[{"document_id": "doc-1", "metadata_list": []}]),
MagicMock(),
current_user,
dataset_id,
)
assert status == 204
assert result == ""