diff --git a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py index 9fae95edcfc..8273e179e2f 100644 --- a/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py +++ b/api/tests/unit_tests/controllers/service_api/dataset/test_dataset_segment.py @@ -16,10 +16,11 @@ Focus on: import inspect import uuid -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import ANY, Mock, patch import pytest from flask import Flask +from sqlalchemy.orm import Session from werkzeug.exceptions import NotFound from controllers.service_api.dataset.segment import ( @@ -36,28 +37,13 @@ from controllers.service_api.dataset.segment import ( ) from core.rag.index_processor.constant.index_type import IndexStructureType from libs.datetime_utils import naive_utc_now -from models.dataset import ChildChunk, Dataset, Document, DocumentSegment +from models.account import Account, Tenant +from models.dataset import ChildChunk, Dataset, Document, DocumentSegment, DocumentSegmentSummary from models.enums import IndexingStatus, SegmentType +from services.api_token_service import CachedApiToken from services.dataset_service import DocumentService, SegmentService -def _session_factory_mock(): - mock_factory = MagicMock() - mock_session = MagicMock() - mock_factory.session = mock_session - - transaction = MagicMock() - transaction.__enter__.return_value = mock_session - transaction.__exit__.return_value = None - mock_factory.get_session_maker.return_value.begin.return_value = transaction - - read_session = MagicMock() - read_session.__enter__.return_value = mock_session - read_session.__exit__.return_value = None - mock_factory.create_session.return_value = read_session - return mock_factory - - def _segment_response_dict(summary: str | None = None): return { "id": "seg-1", @@ -90,6 +76,31 @@ def _segment_response_dict(summary: str | None = None): } +@pytest.fixture +def mock_tenant() -> Tenant: + tenant = Tenant(name="Segment API Tenant") + tenant.id = str(uuid.uuid4()) + return tenant + + +def _account() -> Account: + account = Account(name="Segment API User", email=f"segment-api-{uuid.uuid4()}@example.com") + account.id = str(uuid.uuid4()) + return account + + +def _api_token(tenant_id: str) -> CachedApiToken: + return CachedApiToken( + id=str(uuid.uuid4()), + app_id=None, + tenant_id=tenant_id, + type="dataset", + token="test_token", + last_used_at=None, + created_at=None, + ) + + def _child_chunk() -> ChildChunk: child_chunk = ChildChunk( tenant_id="tenant-1", @@ -110,15 +121,21 @@ def _child_chunk() -> ChildChunk: def _document_for_dataset( dataset: Dataset, document_id: str = "doc-id", doc_form: str = IndexStructureType.PARAGRAPH_INDEX -): - document = Mock() - document.id = document_id - document.dataset_id = dataset.id - document.tenant_id = dataset.tenant_id - document.indexing_status = "completed" - document.enabled = True - document.doc_form = doc_form - return document +) -> Document: + return Document( + id=document_id, + dataset_id=dataset.id, + tenant_id=dataset.tenant_id, + position=1, + data_source_type="upload_file", + batch="batch-1", + name="Document", + created_from="api", + created_by="account-1", + indexing_status="completed", + enabled=True, + doc_form=doc_form, + ) class TestSegmentCreatePayload: @@ -392,7 +409,7 @@ class TestSegmentServiceMockedBehavior: return segment @patch.object(SegmentService, "multi_create_segment") - def test_create_segments_returns_list(self, mock_create, mock_dataset, mock_document): + def test_create_segments_returns_list(self, mock_create, mock_dataset, mock_document, unbound_session: Session): """Test segment creation returns list of segments.""" mock_segments = [ DocumentSegment( @@ -417,13 +434,11 @@ class TestSegmentServiceMockedBehavior: ), ] mock_create.return_value = mock_segments - session = Mock() - result = SegmentService.multi_create_segment( segments=[{"content": "Test"}, {"content": "Test 2"}], document=mock_document, dataset=mock_dataset, - session=session, + session=unbound_session, ) assert result is not None @@ -431,9 +446,19 @@ class TestSegmentServiceMockedBehavior: mock_create.assert_called_once() @patch.object(SegmentService, "get_segments") - def test_get_segments_returns_tuple(self, mock_get, mock_document): + def test_get_segments_returns_tuple(self, mock_get, mock_document, mock_segment, unbound_session: Session): """Test get_segments returns tuple of segments and count.""" - mock_segments = [Mock(), Mock()] + second_segment = DocumentSegment( + tenant_id=mock_segment.tenant_id, + dataset_id=mock_segment.dataset_id, + document_id=mock_segment.document_id, + position=2, + content="Second segment", + word_count=2, + tokens=2, + created_by="account-id", + ) + mock_segments = [mock_segment, second_segment] mock_get.return_value = (mock_segments, 2) segments, count = SegmentService.get_segments( @@ -441,46 +466,43 @@ class TestSegmentServiceMockedBehavior: tenant_id=mock_document.tenant_id, page=1, limit=20, - session=MagicMock(), + session=unbound_session, ) assert len(segments) == 2 assert count == 2 @patch.object(SegmentService, "get_segment_by_id") - def test_get_segment_by_id_returns_segment(self, mock_get, mock_segment): + def test_get_segment_by_id_returns_segment(self, mock_get, mock_segment, unbound_session: Session): """Test get_segment_by_id returns segment.""" mock_get.return_value = mock_segment - session = Mock() - result = SegmentService.get_segment_by_id( segment_id=mock_segment.id, tenant_id=mock_segment.tenant_id, - session=session, + session=unbound_session, ) assert result == mock_segment @patch.object(SegmentService, "get_segment_by_id") - def test_get_segment_by_id_returns_none_when_not_found(self, mock_get): + def test_get_segment_by_id_returns_none_when_not_found(self, mock_get, unbound_session: Session): """Test get_segment_by_id returns None when not found.""" mock_get.return_value = None - session = Mock() - result = SegmentService.get_segment_by_id( segment_id=str(uuid.uuid4()), tenant_id=str(uuid.uuid4()), - session=session, + session=unbound_session, ) assert result is None @patch.object(SegmentService, "delete_segment") - def test_delete_segment_called(self, mock_delete, mock_segment, mock_document, mock_dataset): + def test_delete_segment_called( + self, mock_delete, mock_segment, mock_document, mock_dataset, unbound_session: Session + ): """Test segment deletion is called.""" - session = Mock() - SegmentService.delete_segment(mock_segment, mock_document, mock_dataset, session) - mock_delete.assert_called_once_with(mock_segment, mock_document, mock_dataset, session) + SegmentService.delete_segment(mock_segment, mock_document, mock_dataset, unbound_session) + mock_delete.assert_called_once_with(mock_segment, mock_document, mock_dataset, unbound_session) class TestChildChunkServiceMockedBehavior: @@ -519,7 +541,9 @@ class TestChildChunkServiceMockedBehavior: return chunk @patch.object(SegmentService, "create_child_chunk") - def test_create_child_chunk_returns_chunk(self, mock_create, mock_segment, mock_child_chunk): + def test_create_child_chunk_returns_chunk( + self, mock_create, mock_segment, mock_child_chunk, unbound_session: Session + ): """Test child chunk creation returns chunk.""" mock_create.return_value = mock_child_chunk @@ -528,16 +552,19 @@ class TestChildChunkServiceMockedBehavior: segment=mock_segment, document=Document(), dataset=Dataset(), - session=Mock(), + session=unbound_session, ) assert result == mock_child_chunk @patch.object(SegmentService, "get_child_chunks") - def test_get_child_chunks_returns_paginated_result(self, mock_get, mock_segment): + def test_get_child_chunks_returns_paginated_result(self, mock_get, mock_segment, unbound_session: Session): """Test get_child_chunks returns paginated result.""" mock_pagination = Mock() - mock_pagination.items = [Mock(), Mock()] + first_chunk = _child_chunk() + second_chunk = _child_chunk() + second_chunk.id = "child-2" + mock_pagination.items = [first_chunk, second_chunk] mock_pagination.total = 2 mock_pagination.pages = 1 mock_get.return_value = mock_pagination @@ -548,27 +575,27 @@ class TestChildChunkServiceMockedBehavior: dataset_id=str(uuid.uuid4()), page=1, limit=20, - session=MagicMock(), + session=unbound_session, ) assert len(result.items) == 2 assert result.total == 2 @patch.object(SegmentService, "get_child_chunk_by_id") - def test_get_child_chunk_by_id_returns_chunk(self, mock_get, mock_child_chunk): + def test_get_child_chunk_by_id_returns_chunk(self, mock_get, mock_child_chunk, unbound_session: Session): """Test get_child_chunk_by_id returns chunk.""" mock_get.return_value = mock_child_chunk result = SegmentService.get_child_chunk_by_id( child_chunk_id=mock_child_chunk.id, tenant_id=mock_child_chunk.tenant_id, - session=Mock(), + session=unbound_session, ) assert result == mock_child_chunk @patch.object(SegmentService, "update_child_chunk") - def test_update_child_chunk_returns_updated_chunk(self, mock_update, mock_child_chunk): + def test_update_child_chunk_returns_updated_chunk(self, mock_update, mock_child_chunk, unbound_session: Session): """Test update_child_chunk returns updated chunk.""" updated_chunk = ChildChunk( tenant_id="tenant-id", @@ -597,7 +624,7 @@ class TestChildChunkServiceMockedBehavior: ), document=Document(), dataset=Dataset(), - session=Mock(), + session=unbound_session, ) assert result.content == "Updated content" @@ -977,7 +1004,21 @@ class TestSegmentPagination: # ============================================================================= -class TestSegmentApiGet: +class SQLiteEndpointTest: + session: Session + + @pytest.fixture(autouse=True) + def _use_sqlite_session(self, sqlite_session: Session) -> None: + self.session = sqlite_session + + def _persist_dataset(self, dataset: Dataset, tenant_id: str) -> None: + dataset.tenant_id = tenant_id + dataset.created_by = "account-1" + self.session.add(dataset) + self.session.commit() + + +class TestSegmentApiGet(SQLiteEndpointTest): """Test suite for SegmentApi.get() endpoint. ``get`` has no billing decorators but calls @@ -989,10 +1030,8 @@ class TestSegmentApiGet: @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_list_segments_success( self, - session_factory, mock_account_fn, mock_doc_svc, mock_seg_svc, @@ -1005,8 +1044,8 @@ class TestSegmentApiGet: ): """Test successful segment list retrieval.""" # Arrange - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc_svc.get_document.return_value = _document_for_dataset( mock_dataset, doc_form=IndexStructureType.PARAGRAPH_INDEX ) @@ -1027,15 +1066,14 @@ class TestSegmentApiGet: assert "data" in response assert "total" in response assert response["page"] == 1 - mock_dump_segments.assert_called_once_with([mock_segment], {}, session=session_factory.session) + mock_dump_segments.assert_called_once_with([mock_segment], {}, session=ANY) + assert isinstance(mock_dump_segments.call_args.kwargs["session"], Session) @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) - def test_list_segments_dataset_not_found(self, session_factory, mock_account_fn, app, mock_tenant, mock_dataset): + def test_list_segments_dataset_not_found(self, mock_account_fn, app, mock_tenant, mock_dataset): """Test 404 when dataset not found.""" # Arrange - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = None + mock_account_fn.return_value = (_account(), mock_tenant.id) # Act & Assert with app.test_request_context( @@ -1048,14 +1086,11 @@ class TestSegmentApiGet: @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) - def test_list_segments_document_not_found( - self, session_factory, mock_account_fn, mock_doc_svc, app, mock_tenant, mock_dataset - ): + def test_list_segments_document_not_found(self, mock_account_fn, mock_doc_svc, app, mock_tenant, mock_dataset): """Test 404 when document not found.""" # Arrange - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc_svc.get_document.return_value = None # Act & Assert @@ -1068,7 +1103,7 @@ class TestSegmentApiGet: api.get(tenant_id=mock_tenant.id, dataset_id=mock_dataset.id, document_id="doc-id") -class TestSegmentApiPost: +class TestSegmentApiPost(SQLiteEndpointTest): """Test suite for SegmentApi.post() endpoint. ``post`` is wrapped by ``@cloud_edition_billing_resource_check``, @@ -1082,9 +1117,7 @@ class TestSegmentApiPost: @staticmethod def _setup_billing_mocks(mock_validate_token, mock_feature_svc, tenant_id: str): """Configure mocks to neutralise billing/auth decorators.""" - mock_api_token = Mock() - mock_api_token.tenant_id = tenant_id - mock_validate_token.return_value = mock_api_token + mock_validate_token.return_value = _api_token(tenant_id) mock_features = Mock() mock_features.billing.enabled = False @@ -1104,14 +1137,12 @@ class TestSegmentApiPost: @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_create_segments_success( self, mock_validate_token, mock_feature_svc, - session_factory, mock_account_fn, mock_doc_svc, mock_seg_svc, @@ -1125,10 +1156,10 @@ class TestSegmentApiPost: """Test successful segment creation.""" # Arrange — neutralise billing decorators self._setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant.id) - mock_account_fn.return_value = (Mock(), mock_tenant.id) + mock_account_fn.return_value = (_account(), mock_tenant.id) mock_dataset.indexing_technique = "economy" - session_factory.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc = _document_for_dataset(mock_dataset) mock_doc.indexing_status = "completed" @@ -1161,18 +1192,17 @@ class TestSegmentApiPost: assert status == 200 assert "data" in response assert "doc_form" in response - mock_dump_segments.assert_called_once_with([mock_segment], {}, session=session_factory.session) + mock_dump_segments.assert_called_once_with([mock_segment], {}, session=ANY) + assert isinstance(mock_dump_segments.call_args.kwargs["session"], Session) @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_create_segments_missing_segments( self, mock_validate_token, mock_feature_svc, - session_factory, mock_account_fn, mock_doc_svc, app: Flask, @@ -1182,10 +1212,10 @@ class TestSegmentApiPost: """Test 400 error when segments field is missing.""" # Arrange — neutralise billing decorators self._setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant.id) - mock_account_fn.return_value = (Mock(), mock_tenant.id) + mock_account_fn.return_value = (_account(), mock_tenant.id) mock_dataset.indexing_technique = "economy" - session_factory.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc = _document_for_dataset(mock_dataset) mock_doc.indexing_status = "completed" @@ -1212,14 +1242,12 @@ class TestSegmentApiPost: @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_create_segments_document_not_completed( self, mock_validate_token, mock_feature_svc, - session_factory, mock_account_fn, mock_doc_svc, app: Flask, @@ -1229,9 +1257,9 @@ class TestSegmentApiPost: """Test 404 when document indexing is not completed.""" # Arrange — neutralise billing decorators self._setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant.id) - mock_account_fn.return_value = (Mock(), mock_tenant.id) + mock_account_fn.return_value = (_account(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc = _document_for_dataset(mock_dataset) mock_doc.indexing_status = "indexing" # Not completed @@ -1253,7 +1281,7 @@ class TestSegmentApiPost: ) -class TestDatasetSegmentApiDelete: +class TestDatasetSegmentApiDelete(SQLiteEndpointTest): """Test suite for DatasetSegmentApi.delete() endpoint. ``delete`` is wrapped by ``@cloud_edition_billing_rate_limit_check`` @@ -1265,10 +1293,8 @@ class TestDatasetSegmentApiDelete: @patch("controllers.service_api.dataset.segment.DatasetService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_delete_segment_success( self, - session_factory, mock_account_fn, mock_doc_svc, mock_dataset_svc, @@ -1280,8 +1306,8 @@ class TestDatasetSegmentApiDelete: ): """Test successful segment deletion.""" # Arrange - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_dataset_svc.check_dataset_model_setting.return_value = None mock_doc = _document_for_dataset(mock_dataset) @@ -1299,7 +1325,7 @@ class TestDatasetSegmentApiDelete: delete = inspect.unwrap(api.delete) response = delete( api, - session_factory.session, + self.session, tenant_id=mock_tenant.id, dataset_id=mock_dataset.id, document_id="doc-id", @@ -1308,17 +1334,13 @@ class TestDatasetSegmentApiDelete: # Assert assert response == ("", 204) - mock_seg_svc.delete_segment.assert_called_once_with( - mock_segment, mock_doc, mock_dataset, session_factory.session - ) + mock_seg_svc.delete_segment.assert_called_once_with(mock_segment, mock_doc, mock_dataset, self.session) @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_delete_segment_not_found( self, - session_factory, mock_account_fn, mock_doc_svc, mock_seg_svc, @@ -1328,8 +1350,8 @@ class TestDatasetSegmentApiDelete: ): """Test 404 when segment not found.""" # Arrange - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc = _document_for_dataset(mock_dataset) mock_doc.indexing_status = "completed" @@ -1349,7 +1371,7 @@ class TestDatasetSegmentApiDelete: with pytest.raises(NotFound): delete( api, - session_factory.session, + self.session, tenant_id=mock_tenant.id, dataset_id=mock_dataset.id, document_id="doc-id", @@ -1359,10 +1381,8 @@ class TestDatasetSegmentApiDelete: @patch("controllers.service_api.dataset.segment.DatasetService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_delete_segment_dataset_not_found( self, - session_factory, mock_account_fn, mock_doc_svc, mock_dataset_svc, @@ -1372,8 +1392,7 @@ class TestDatasetSegmentApiDelete: ): """Test 404 when dataset not found for delete.""" # Arrange - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = None + mock_account_fn.return_value = (_account(), mock_tenant.id) # Act & Assert with app.test_request_context( @@ -1385,7 +1404,7 @@ class TestDatasetSegmentApiDelete: with pytest.raises(NotFound): delete( api, - session_factory.session, + self.session, tenant_id=mock_tenant.id, dataset_id=mock_dataset.id, document_id="doc-id", @@ -1395,10 +1414,8 @@ class TestDatasetSegmentApiDelete: @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.DatasetService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_delete_segment_document_not_found( self, - session_factory, mock_account_fn, mock_dataset_svc, mock_doc_svc, @@ -1408,8 +1425,8 @@ class TestDatasetSegmentApiDelete: ): """Test 404 when document not found for delete.""" # Arrange - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_dataset_svc.check_dataset_model_setting.return_value = None mock_doc_svc.get_document.return_value = None @@ -1423,7 +1440,7 @@ class TestDatasetSegmentApiDelete: with pytest.raises(NotFound): delete( api, - session_factory.session, + self.session, tenant_id=mock_tenant.id, dataset_id=mock_dataset.id, document_id="doc-id", @@ -1431,7 +1448,7 @@ class TestDatasetSegmentApiDelete: ) -class TestDatasetSegmentApiUpdate: +class TestDatasetSegmentApiUpdate(SQLiteEndpointTest): """Test suite for DatasetSegmentApi.post() (update segment) endpoint. ``post`` is wrapped by ``@cloud_edition_billing_resource_check`` and @@ -1444,9 +1461,7 @@ class TestDatasetSegmentApiUpdate: @staticmethod def _setup_billing_mocks(mock_validate_token, mock_feature_svc, tenant_id: str): """Configure mocks to neutralise billing/auth decorators.""" - mock_api_token = Mock() - mock_api_token.tenant_id = tenant_id - mock_validate_token.return_value = mock_api_token + mock_validate_token.return_value = _api_token(tenant_id) mock_features = Mock() mock_features.billing.enabled = False mock_feature_svc.get_features.return_value = mock_features @@ -1464,14 +1479,12 @@ class TestDatasetSegmentApiUpdate: @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.DatasetService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_update_segment_success( self, mock_validate_token, mock_feature_svc, - session_factory, mock_account_fn, mock_dataset_svc, mock_doc_svc, @@ -1485,15 +1498,24 @@ class TestDatasetSegmentApiUpdate: ): """Test successful segment update.""" self._setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant.id) - mock_account_fn.return_value = (Mock(), mock_tenant.id) + mock_account_fn.return_value = (_account(), mock_tenant.id) mock_dataset.indexing_technique = "economy" - session_factory.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset, mock_tenant.id) mock_dataset_svc.check_dataset_model_setting.return_value = None mock_doc_svc.get_document.return_value = _document_for_dataset( mock_dataset, doc_form=IndexStructureType.PARAGRAPH_INDEX ) mock_seg_svc.get_segment_by_ref.return_value = mock_segment - updated = Mock() + updated = DocumentSegment( + tenant_id=mock_tenant.id, + dataset_id=mock_dataset.id, + document_id="doc-id", + position=1, + content="updated content", + word_count=2, + tokens=2, + created_by="account-1", + ) updated.id = "updated-seg" mock_seg_svc.update_segment.return_value = updated mock_get_summary.return_value = None @@ -1516,19 +1538,18 @@ class TestDatasetSegmentApiUpdate: assert status == 200 assert "data" in response mock_seg_svc.update_segment.assert_called_once() - mock_dump_segment.assert_called_once_with(updated, None, session=session_factory.session) + mock_dump_segment.assert_called_once_with(updated, None, session=ANY) + assert isinstance(mock_dump_segment.call_args.kwargs["session"], Session) @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.DatasetService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_update_segment_dataset_not_found( self, mock_validate_token, mock_feature_svc, - session_factory, mock_account_fn, mock_dataset_svc, mock_doc_svc, @@ -1538,8 +1559,7 @@ class TestDatasetSegmentApiUpdate: ): """Test 404 when dataset not found for update.""" self._setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant.id) - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = None + mock_account_fn.return_value = (_account(), mock_tenant.id) with app.test_request_context( f"/datasets/{mock_dataset.id}/documents/doc-id/segments/seg-id", @@ -1560,14 +1580,12 @@ class TestDatasetSegmentApiUpdate: @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.DatasetService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_update_segment_not_found( self, mock_validate_token, mock_feature_svc, - session_factory, mock_account_fn, mock_dataset_svc, mock_doc_svc, @@ -1578,9 +1596,9 @@ class TestDatasetSegmentApiUpdate: ): """Test 404 when segment not found for update.""" self._setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant.id) - mock_account_fn.return_value = (Mock(), mock_tenant.id) + mock_account_fn.return_value = (_account(), mock_tenant.id) mock_dataset.indexing_technique = "economy" - session_factory.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset, mock_tenant.id) mock_dataset_svc.check_dataset_model_setting.return_value = None mock_doc_svc.get_document.return_value = _document_for_dataset(mock_dataset) mock_seg_svc.get_segment_by_ref.return_value = None @@ -1601,7 +1619,7 @@ class TestDatasetSegmentApiUpdate: ) -class TestDatasetSegmentApiGetSingle: +class TestDatasetSegmentApiGetSingle(SQLiteEndpointTest): """Test suite for DatasetSegmentApi.get() (single segment) endpoint. ``get`` has no billing decorators but calls @@ -1614,10 +1632,8 @@ class TestDatasetSegmentApiGetSingle: @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.DatasetService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_get_single_segment_success( self, - session_factory, mock_account_fn, mock_dataset_svc, mock_doc_svc, @@ -1630,8 +1646,8 @@ class TestDatasetSegmentApiGetSingle: mock_segment, ): """Test successful single segment retrieval.""" - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_dataset_svc.check_dataset_model_setting.return_value = None mock_doc = _document_for_dataset(mock_dataset, doc_form=IndexStructureType.PARAGRAPH_INDEX) mock_doc_svc.get_document.return_value = mock_doc @@ -1654,7 +1670,8 @@ class TestDatasetSegmentApiGetSingle: assert status == 200 assert "data" in response assert response["doc_form"] == IndexStructureType.PARAGRAPH_INDEX - mock_dump_segment.assert_called_once_with(mock_segment, None, session=session_factory.session) + mock_dump_segment.assert_called_once_with(mock_segment, None, session=ANY) + assert isinstance(mock_dump_segment.call_args.kwargs["session"], Session) @patch("controllers.service_api.dataset.segment.segment_response_with_summary") @patch("controllers.service_api.dataset.segment.SummaryIndexService.get_segment_summary") @@ -1662,10 +1679,8 @@ class TestDatasetSegmentApiGetSingle: @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.DatasetService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_get_single_segment_includes_summary( self, - session_factory, mock_account_fn, mock_dataset_svc, mock_doc_svc, @@ -1678,13 +1693,18 @@ class TestDatasetSegmentApiGetSingle: mock_segment, ): """Test that single segment response includes summary content from SummaryIndexService.""" - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_dataset_svc.check_dataset_model_setting.return_value = None mock_doc = _document_for_dataset(mock_dataset, doc_form=IndexStructureType.PARAGRAPH_INDEX) mock_doc_svc.get_document.return_value = mock_doc mock_seg_svc.get_segment_by_ref.return_value = mock_segment - mock_summary_record = Mock(summary_content="This is the segment summary") + mock_summary_record = DocumentSegmentSummary( + dataset_id=mock_dataset.id, + document_id="doc-id", + chunk_id=mock_segment.id, + summary_content="This is the segment summary", + ) mock_get_summary.return_value = mock_summary_record mock_dump_segment.return_value = _segment_response_dict("This is the segment summary") @@ -1705,22 +1725,20 @@ class TestDatasetSegmentApiGetSingle: mock_dump_segment.assert_called_once_with( mock_segment, "This is the segment summary", - session=session_factory.session, + session=ANY, ) + assert isinstance(mock_dump_segment.call_args.kwargs["session"], Session) @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_get_single_segment_dataset_not_found( self, - session_factory, mock_account_fn, app: Flask, mock_tenant, mock_dataset, ): """Test 404 when dataset not found.""" - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = None + mock_account_fn.return_value = (_account(), mock_tenant.id) with app.test_request_context( f"/datasets/{mock_dataset.id}/documents/doc-id/segments/seg-id", @@ -1738,10 +1756,8 @@ class TestDatasetSegmentApiGetSingle: @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.DatasetService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_get_single_segment_document_not_found( self, - session_factory, mock_account_fn, mock_dataset_svc, mock_doc_svc, @@ -1750,8 +1766,8 @@ class TestDatasetSegmentApiGetSingle: mock_dataset, ): """Test 404 when document not found.""" - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_dataset_svc.check_dataset_model_setting.return_value = None mock_doc_svc.get_document.return_value = None @@ -1772,10 +1788,8 @@ class TestDatasetSegmentApiGetSingle: @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.DatasetService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_get_single_segment_segment_not_found( self, - session_factory, mock_account_fn, mock_dataset_svc, mock_doc_svc, @@ -1785,8 +1799,8 @@ class TestDatasetSegmentApiGetSingle: mock_dataset, ): """Test 404 when segment not found.""" - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_dataset_svc.check_dataset_model_setting.return_value = None mock_doc_svc.get_document.return_value = _document_for_dataset(mock_dataset) mock_seg_svc.get_segment_by_ref.return_value = None @@ -1805,7 +1819,7 @@ class TestDatasetSegmentApiGetSingle: ) -class TestChildChunkApiGet: +class TestChildChunkApiGet(SQLiteEndpointTest): """Test suite for ChildChunkApi.get() endpoint. ``get`` has no billing decorators but calls @@ -1815,22 +1829,21 @@ class TestChildChunkApiGet: @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_list_child_chunks_success( self, - session_factory, mock_account_fn, mock_doc_svc, mock_seg_svc, app: Flask, mock_tenant, mock_dataset, + mock_segment, ): """Test successful child chunk list retrieval.""" - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc_svc.get_document.return_value = _document_for_dataset(mock_dataset) - mock_seg_svc.get_segment_by_ref.return_value = Mock() + mock_seg_svc.get_segment_by_ref.return_value = mock_segment mock_pagination = Mock() mock_pagination.items = [_child_chunk(), _child_chunk()] @@ -1855,18 +1868,15 @@ class TestChildChunkApiGet: assert response["page"] == 1 @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_list_child_chunks_dataset_not_found( self, - session_factory, mock_account_fn, app: Flask, mock_tenant, mock_dataset, ): """Test 404 when dataset not found.""" - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = None + mock_account_fn.return_value = (_account(), mock_tenant.id) with app.test_request_context( f"/datasets/{mock_dataset.id}/documents/doc-id/segments/seg-id/child_chunks", @@ -1883,10 +1893,8 @@ class TestChildChunkApiGet: @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_list_child_chunks_document_not_found( self, - session_factory, mock_account_fn, mock_doc_svc, app: Flask, @@ -1894,8 +1902,8 @@ class TestChildChunkApiGet: mock_dataset, ): """Test 404 when document not found.""" - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc_svc.get_document.return_value = None with app.test_request_context( @@ -1914,10 +1922,8 @@ class TestChildChunkApiGet: @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_list_child_chunks_segment_not_found( self, - session_factory, mock_account_fn, mock_doc_svc, mock_seg_svc, @@ -1926,8 +1932,8 @@ class TestChildChunkApiGet: mock_dataset, ): """Test 404 when segment not found.""" - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc_svc.get_document.return_value = _document_for_dataset(mock_dataset) mock_seg_svc.get_segment_by_ref.return_value = None @@ -1945,7 +1951,7 @@ class TestChildChunkApiGet: ) -class TestChildChunkApiPost: +class TestChildChunkApiPost(SQLiteEndpointTest): """Test suite for ChildChunkApi.post() endpoint. ``post`` has billing decorators; we patch ``validate_and_get_api_token`` @@ -1954,9 +1960,7 @@ class TestChildChunkApiPost: @staticmethod def _setup_billing_mocks(mock_validate_token, mock_feature_svc, tenant_id: str): - mock_api_token = Mock() - mock_api_token.tenant_id = tenant_id - mock_validate_token.return_value = mock_api_token + mock_validate_token.return_value = _api_token(tenant_id) mock_features = Mock() mock_features.billing.enabled = False mock_feature_svc.get_features.return_value = mock_features @@ -1971,28 +1975,27 @@ class TestChildChunkApiPost: @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_create_child_chunk_success( self, mock_validate_token, mock_feature_svc, - session_factory, mock_account_fn, mock_doc_svc, mock_seg_svc, app: Flask, mock_tenant, mock_dataset, + mock_segment, ): """Test successful child chunk creation.""" self._setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant.id) - mock_account_fn.return_value = (Mock(), mock_tenant.id) + mock_account_fn.return_value = (_account(), mock_tenant.id) mock_dataset.indexing_technique = "economy" - session_factory.session.scalar.return_value = mock_dataset + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc_svc.get_document.return_value = _document_for_dataset(mock_dataset) - mock_seg_svc.get_segment_by_ref.return_value = Mock() + mock_seg_svc.get_segment_by_ref.return_value = mock_segment mock_child = _child_chunk() mock_seg_svc.create_child_chunk.return_value = mock_child @@ -2014,14 +2017,12 @@ class TestChildChunkApiPost: assert "data" in response @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_create_child_chunk_dataset_not_found( self, mock_validate_token, mock_feature_svc, - session_factory, mock_account_fn, app: Flask, mock_tenant, @@ -2029,8 +2030,7 @@ class TestChildChunkApiPost: ): """Test 404 when dataset not found.""" self._setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant.id) - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = None + mock_account_fn.return_value = (_account(), mock_tenant.id) with app.test_request_context( f"/datasets/{mock_dataset.id}/documents/doc-id/segments/seg-id/child_chunks", @@ -2050,14 +2050,12 @@ class TestChildChunkApiPost: @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) @patch("controllers.service_api.wraps.FeatureService") @patch("controllers.service_api.wraps.validate_and_get_api_token") def test_create_child_chunk_segment_not_found( self, mock_validate_token, mock_feature_svc, - session_factory, mock_account_fn, mock_doc_svc, mock_seg_svc, @@ -2067,8 +2065,8 @@ class TestChildChunkApiPost: ): """Test 404 when segment not found.""" self._setup_billing_mocks(mock_validate_token, mock_feature_svc, mock_tenant.id) - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc_svc.get_document.return_value = _document_for_dataset(mock_dataset) mock_seg_svc.get_segment_by_ref.return_value = None @@ -2088,7 +2086,7 @@ class TestChildChunkApiPost: ) -class TestDatasetChildChunkApiDelete: +class TestDatasetChildChunkApiDelete(SQLiteEndpointTest): """Test suite for DatasetChildChunkApi.delete() endpoint. ``delete`` is wrapped by ``@cloud_edition_billing_knowledge_limit_check`` @@ -2100,34 +2098,33 @@ class TestDatasetChildChunkApiDelete: @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_delete_child_chunk_success( self, - session_factory, mock_account_fn, mock_doc_svc, mock_seg_svc, app: Flask, mock_tenant, mock_dataset, + mock_segment, + mock_child_chunk, ): """Test successful child chunk deletion.""" - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc = _document_for_dataset(mock_dataset) mock_doc_svc.get_document.return_value = mock_doc segment_id = str(uuid.uuid4()) - mock_segment = Mock() mock_segment.id = segment_id mock_segment.document_id = "doc-id" mock_seg_svc.get_segment_by_ref.return_value = mock_segment child_chunk_id = str(uuid.uuid4()) - mock_child = Mock() - mock_child.segment_id = segment_id - mock_seg_svc.get_child_chunk_by_segment_ref.return_value = mock_child + mock_child_chunk.id = child_chunk_id + mock_child_chunk.segment_id = segment_id + mock_seg_svc.get_child_chunk_by_segment_ref.return_value = mock_child_chunk mock_seg_svc.delete_child_chunk.return_value = None with app.test_request_context( @@ -2138,7 +2135,7 @@ class TestDatasetChildChunkApiDelete: delete = inspect.unwrap(api.delete) response = delete( api, - session_factory.session, + self.session, tenant_id=mock_tenant.id, dataset_id=mock_dataset.id, document_id="doc-id", @@ -2152,24 +2149,22 @@ class TestDatasetChildChunkApiDelete: @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_delete_child_chunk_not_found( self, - session_factory, mock_account_fn, mock_doc_svc, mock_seg_svc, app: Flask, mock_tenant, mock_dataset, + mock_segment, ): """Test 404 when child chunk not found.""" - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc_svc.get_document.return_value = _document_for_dataset(mock_dataset) segment_id = str(uuid.uuid4()) - mock_segment = Mock() mock_segment.id = segment_id mock_segment.document_id = "doc-id" mock_seg_svc.get_segment_by_ref.return_value = mock_segment @@ -2184,7 +2179,7 @@ class TestDatasetChildChunkApiDelete: with pytest.raises(NotFound): delete( api, - session_factory.session, + self.session, tenant_id=mock_tenant.id, dataset_id=mock_dataset.id, document_id="doc-id", @@ -2195,10 +2190,8 @@ class TestDatasetChildChunkApiDelete: @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_delete_child_chunk_segment_document_mismatch( self, - session_factory, mock_account_fn, mock_doc_svc, mock_seg_svc, @@ -2207,8 +2200,8 @@ class TestDatasetChildChunkApiDelete: mock_dataset, ): """Test 404 when segment does not belong to the document.""" - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc_svc.get_document.return_value = _document_for_dataset(mock_dataset) segment_id = str(uuid.uuid4()) @@ -2223,7 +2216,7 @@ class TestDatasetChildChunkApiDelete: with pytest.raises(NotFound): delete( api, - session_factory.session, + self.session, tenant_id=mock_tenant.id, dataset_id=mock_dataset.id, document_id="doc-id", @@ -2234,24 +2227,22 @@ class TestDatasetChildChunkApiDelete: @patch("controllers.service_api.dataset.segment.SegmentService") @patch("controllers.service_api.dataset.segment.DocumentService") @patch("controllers.service_api.dataset.segment.current_account_with_tenant") - @patch("controllers.common.session.session_factory", new_callable=_session_factory_mock) def test_delete_child_chunk_wrong_segment( self, - session_factory, mock_account_fn, mock_doc_svc, mock_seg_svc, app: Flask, mock_tenant, mock_dataset, + mock_segment, ): """Test 404 when child chunk does not belong to the segment.""" - mock_account_fn.return_value = (Mock(), mock_tenant.id) - session_factory.session.scalar.return_value = mock_dataset + mock_account_fn.return_value = (_account(), mock_tenant.id) + self._persist_dataset(mock_dataset, mock_tenant.id) mock_doc_svc.get_document.return_value = _document_for_dataset(mock_dataset) segment_id = str(uuid.uuid4()) - mock_segment = Mock() mock_segment.id = segment_id mock_segment.document_id = "doc-id" mock_seg_svc.get_segment_by_ref.return_value = mock_segment @@ -2267,7 +2258,7 @@ class TestDatasetChildChunkApiDelete: with pytest.raises(NotFound): delete( api, - session_factory.session, + self.session, tenant_id=mock_tenant.id, dataset_id=mock_dataset.id, document_id="doc-id",