refactor: select in console datasets document controller (#34019)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Renzo
2026-03-24 13:57:38 +01:00
committed by GitHub
parent 542c1a14e0
commit e3c1112b15
4 changed files with 82 additions and 96 deletions

View File

@ -140,8 +140,8 @@ class TestDatasetDocumentListApi:
return_value=pagination,
),
patch(
"controllers.console.datasets.datasets_document.db.session.query",
return_value=MagicMock(where=lambda *a, **k: MagicMock(count=count_mock)),
"controllers.console.datasets.datasets_document.db.session.scalar",
return_value=2,
),
patch(
"controllers.console.datasets.datasets_document.DocumentService.enrich_documents_with_summary_index_status",
@ -700,10 +700,8 @@ class TestDocumentPipelineExecutionLogApi:
return_value=MagicMock(),
),
patch(
"controllers.console.datasets.datasets_document.db.session.query",
return_value=MagicMock(
filter_by=lambda **k: MagicMock(order_by=lambda *a: MagicMock(first=lambda: log))
),
"controllers.console.datasets.datasets_document.db.session.scalar",
return_value=log,
),
):
response, status = method(api, "ds-1", "doc-1")
@ -827,15 +825,12 @@ class TestDocumentIndexingEstimateApi:
dataset_process_rule=None,
)
query_mock = MagicMock()
query_mock.where.return_value.first.return_value = None
with (
app.test_request_context("/"),
patch.object(api, "get_document", return_value=document),
patch(
"controllers.console.datasets.datasets_document.db.session.query",
return_value=query_mock,
"controllers.console.datasets.datasets_document.db.session.scalar",
return_value=None,
),
):
with pytest.raises(NotFound):
@ -863,10 +858,8 @@ class TestDocumentIndexingEstimateApi:
app.test_request_context("/"),
patch.object(api, "get_document", return_value=document),
patch(
"controllers.console.datasets.datasets_document.db.session.query",
return_value=MagicMock(
where=MagicMock(return_value=MagicMock(first=MagicMock(return_value=upload_file)))
),
"controllers.console.datasets.datasets_document.db.session.scalar",
return_value=upload_file,
),
patch(
"controllers.console.datasets.datasets_document.ExtractSetting",
@ -1239,12 +1232,8 @@ class TestDocumentPermissionCases:
return_value=None,
),
patch(
"controllers.console.datasets.datasets_document.db.session.query",
return_value=MagicMock(
where=lambda *a: MagicMock(
order_by=lambda *b: MagicMock(limit=lambda n: MagicMock(one_or_none=lambda: process_rule))
)
),
"controllers.console.datasets.datasets_document.db.session.scalar",
return_value=process_rule,
),
):
result = method(api)
@ -1364,8 +1353,8 @@ class TestDocumentIndexingEdgeCases:
app.test_request_context("/"),
patch.object(api, "get_document", return_value=document),
patch(
"controllers.console.datasets.datasets_document.db.session.query",
return_value=MagicMock(where=lambda *a: MagicMock(first=lambda: upload_file)),
"controllers.console.datasets.datasets_document.db.session.scalar",
return_value=upload_file,
),
patch(
"controllers.console.datasets.datasets_document.ExtractSetting",

View File

@ -26,12 +26,9 @@ class TestGetRagPipeline:
return_value=(Mock(), "tenant-1"),
)
mock_query = Mock()
mock_query.where.return_value.first.return_value = None
mocker.patch(
"controllers.console.datasets.wraps.db.session.query",
return_value=mock_query,
"controllers.console.datasets.wraps.db.session.scalar",
return_value=None,
)
with pytest.raises(PipelineNotFoundError):
@ -51,12 +48,9 @@ class TestGetRagPipeline:
return_value=(Mock(), "tenant-1"),
)
mock_query = Mock()
mock_query.where.return_value.first.return_value = pipeline
mocker.patch(
"controllers.console.datasets.wraps.db.session.query",
return_value=mock_query,
"controllers.console.datasets.wraps.db.session.scalar",
return_value=pipeline,
)
result = dummy_view(pipeline_id="pipeline-1")
@ -76,12 +70,9 @@ class TestGetRagPipeline:
return_value=(Mock(), "tenant-1"),
)
mock_query = Mock()
mock_query.where.return_value.first.return_value = pipeline
mocker.patch(
"controllers.console.datasets.wraps.db.session.query",
return_value=mock_query,
"controllers.console.datasets.wraps.db.session.scalar",
return_value=pipeline,
)
result = dummy_view(pipeline_id="pipeline-1")
@ -100,18 +91,15 @@ class TestGetRagPipeline:
return_value=(Mock(), "tenant-1"),
)
def where_side_effect(*args, **kwargs):
assert args[0].right.value == "123"
return Mock(first=lambda: pipeline)
mock_query = Mock()
mock_query.where.side_effect = where_side_effect
mocker.patch(
"controllers.console.datasets.wraps.db.session.query",
return_value=mock_query,
mock_scalar = mocker.patch(
"controllers.console.datasets.wraps.db.session.scalar",
return_value=pipeline,
)
result = dummy_view(pipeline_id=123)
assert result is pipeline
# Verify the pipeline_id was cast to string in the where clause
stmt = mock_scalar.call_args[0][0]
where_clauses = stmt.whereclause.clauses
assert where_clauses[0].right.value == "123"