refactor: select in 10 service files (#34373)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Asuka Minato <i@asukaminato.eu.org>
2026-04-05 09:49:25 +08:00 · 2026-04-01 10:03:49 +02:00
parent b23ea0397a
commit 31f7752ba9
14 changed files with 147 additions and 180 deletions
--- a/api/tests/test_containers_integration_tests/services/test_feedback_service.py
+++ b/api/tests/test_containers_integration_tests/services/test_feedback_service.py
@@ -99,7 +99,7 @@ class TestFeedbackService:
            )
        ]

-        mock_db_session.query.return_value = mock_query
+        mock_db_session.execute.return_value = mock_query

        # Test CSV export
        result = FeedbackService.export_feedbacks(app_id=sample_data["app"].id, format_type="csv")
@@ -138,7 +138,7 @@ class TestFeedbackService:
            )
        ]

-        mock_db_session.query.return_value = mock_query
+        mock_db_session.execute.return_value = mock_query

        # Test JSON export
        result = FeedbackService.export_feedbacks(app_id=sample_data["app"].id, format_type="json")
@@ -175,7 +175,7 @@ class TestFeedbackService:
            )
        ]

-        mock_db_session.query.return_value = mock_query
+        mock_db_session.execute.return_value = mock_query

        # Test with filters
        result = FeedbackService.export_feedbacks(
@@ -188,11 +188,8 @@ class TestFeedbackService:
            format_type="csv",
        )

-        # Verify filters were applied
-        assert mock_query.filter.called
-        filter_calls = mock_query.filter.call_args_list
-        # At least three filter invocations are expected (source, rating, comment)
-        assert len(filter_calls) >= 3
+        # Verify query was executed (filters are baked into the select statement)
+        assert mock_db_session.execute.called

    def test_export_feedbacks_no_data(self, mock_db_session, sample_data):
        """Test exporting feedback when no data exists."""
@@ -206,7 +203,7 @@ class TestFeedbackService:
        mock_query.order_by.return_value = mock_query
        mock_query.all.return_value = []

-        mock_db_session.query.return_value = mock_query
+        mock_db_session.execute.return_value = mock_query

        result = FeedbackService.export_feedbacks(app_id=sample_data["app"].id, format_type="csv")

@@ -271,7 +268,7 @@ class TestFeedbackService:
            )
        ]

-        mock_db_session.query.return_value = mock_query
+        mock_db_session.execute.return_value = mock_query

        # Test export
        result = FeedbackService.export_feedbacks(app_id=sample_data["app"].id, format_type="json")
@@ -329,7 +326,7 @@ class TestFeedbackService:
            )
        ]

-        mock_db_session.query.return_value = mock_query
+        mock_db_session.execute.return_value = mock_query

        # Test export
        result = FeedbackService.export_feedbacks(app_id=sample_data["app"].id, format_type="csv")
@@ -367,7 +364,7 @@ class TestFeedbackService:
            ),
        ]

-        mock_db_session.query.return_value = mock_query
+        mock_db_session.execute.return_value = mock_query

        # Test export
        result = FeedbackService.export_feedbacks(app_id=sample_data["app"].id, format_type="json")
--- a/api/tests/unit_tests/services/test_vector_service.py
+++ b/api/tests/unit_tests/services/test_vector_service.py
@@ -77,22 +77,12 @@ def _make_segment(
 def _mock_db_session_for_update_multimodel(*, upload_files: list[_UploadFileStub] | None) -> MagicMock:
    session = MagicMock(name="session")

-    binding_query = MagicMock(name="binding_query")
-    binding_query.where.return_value = binding_query
-    binding_query.delete.return_value = 1
+    # db.session.execute() is used for delete(SegmentAttachmentBinding).where(...)
+    session.execute = MagicMock(name="execute")

-    upload_query = MagicMock(name="upload_query")
-    upload_query.where.return_value = upload_query
-    upload_query.all.return_value = upload_files or []
+    # db.session.scalars(select(UploadFile).where(...)).all() returns upload files
+    session.scalars.return_value.all.return_value = upload_files or []

-    def query_side_effect(model: object) -> MagicMock:
-        if model is vector_service_module.SegmentAttachmentBinding:
-            return binding_query
-        if model is vector_service_module.UploadFile:
-            return upload_query
-        return MagicMock(name=f"query({model})")
-
-    session.query.side_effect = query_side_effect
    db_mock = MagicMock(name="db")
    db_mock.session = session
    return db_mock
@@ -165,22 +155,15 @@ def _mock_parent_child_queries(
 ) -> MagicMock:
    session = MagicMock(name="session")

-    doc_query = MagicMock(name="doc_query")
-    doc_query.filter_by.return_value = doc_query
-    doc_query.first.return_value = dataset_document
+    get_dispatch: dict[object, object | None] = {
+        vector_service_module.DatasetDocument: dataset_document,
+        vector_service_module.DatasetProcessRule: processing_rule,
+    }

-    rule_query = MagicMock(name="rule_query")
-    rule_query.where.return_value = rule_query
-    rule_query.first.return_value = processing_rule
+    def get_side_effect(model: object, pk: object) -> object | None:
+        return get_dispatch.get(model)

-    def query_side_effect(model: object) -> MagicMock:
-        if model is vector_service_module.DatasetDocument:
-            return doc_query
-        if model is vector_service_module.DatasetProcessRule:
-            return rule_query
-        return MagicMock(name=f"query({model})")
-
-    session.query.side_effect = query_side_effect
+    session.get.side_effect = get_side_effect
    db_mock = MagicMock(name="db")
    db_mock.session = session
    return db_mock
@@ -609,7 +592,7 @@ def test_update_multimodel_vector_deletes_bindings_and_commits_on_empty_new_ids(

    vector_cls.assert_called_once_with(dataset=dataset)
    vector_instance.delete_by_ids.assert_called_once_with(["old-1", "old-2"])
-    db_mock.session.query.assert_called_once_with(vector_service_module.SegmentAttachmentBinding)
+    db_mock.session.execute.assert_called_once()
    db_mock.session.commit.assert_called_once()
    db_mock.session.add_all.assert_not_called()
    vector_instance.add_texts.assert_not_called()
@@ -644,6 +627,8 @@ def test_update_multimodel_vector_adds_bindings_and_vectors_and_skips_missing_up

    binding_ctor = MagicMock(side_effect=lambda **kwargs: kwargs)
    monkeypatch.setattr(vector_service_module, "SegmentAttachmentBinding", binding_ctor)
+    monkeypatch.setattr(vector_service_module, "delete", MagicMock())
+    monkeypatch.setattr(vector_service_module, "select", MagicMock())

    logger_mock = MagicMock()
    monkeypatch.setattr(vector_service_module, "logger", logger_mock)
@@ -677,6 +662,8 @@ def test_update_multimodel_vector_updates_bindings_without_multimodal_vector_ops
    monkeypatch.setattr(
        vector_service_module, "SegmentAttachmentBinding", MagicMock(side_effect=lambda **kwargs: kwargs)
    )
+    monkeypatch.setattr(vector_service_module, "delete", MagicMock())
+    monkeypatch.setattr(vector_service_module, "select", MagicMock())

    VectorService.update_multimodel_vector(segment=segment, attachment_ids=["file-1"], dataset=dataset)

@@ -698,6 +685,8 @@ def test_update_multimodel_vector_rolls_back_and_reraises_on_error(monkeypatch:
    monkeypatch.setattr(
        vector_service_module, "SegmentAttachmentBinding", MagicMock(side_effect=lambda **kwargs: kwargs)
    )
+    monkeypatch.setattr(vector_service_module, "delete", MagicMock())
+    monkeypatch.setattr(vector_service_module, "select", MagicMock())

    logger_mock = MagicMock()
    monkeypatch.setattr(vector_service_module, "logger", logger_mock)
--- a/api/tests/unit_tests/services/test_workflow_service.py
+++ b/api/tests/unit_tests/services/test_workflow_service.py
@@ -268,7 +268,7 @@ class TestWorkflowService:
        Provides mock implementations of:
        - session.add(): Adding new records
        - session.commit(): Committing transactions
-        - session.query(): Querying database
+        - session.scalar(): Scalar queries
        - session.execute(): Executing SQL statements
        """
        with patch("services.workflow_service.db") as mock_db:
@@ -276,7 +276,7 @@ class TestWorkflowService:
            mock_db.session = mock_session
            mock_session.add = MagicMock()
            mock_session.commit = MagicMock()
-            mock_session.query = MagicMock()
+            mock_session.scalar = MagicMock()
            mock_session.execute = MagicMock()
            yield mock_db

@@ -338,10 +338,8 @@ class TestWorkflowService:
        app = TestWorkflowAssociatedDataFactory.create_app_mock()
        mock_workflow = TestWorkflowAssociatedDataFactory.create_workflow_mock()

-        # Mock database query
-        mock_query = MagicMock()
-        mock_db_session.session.query.return_value = mock_query
-        mock_query.where.return_value.first.return_value = mock_workflow
+        # Mock db.session.scalar() used by get_draft_workflow
+        mock_db_session.session.scalar.return_value = mock_workflow

        result = workflow_service.get_draft_workflow(app)

@@ -351,10 +349,8 @@ class TestWorkflowService:
        """Test get_draft_workflow returns None when no draft exists."""
        app = TestWorkflowAssociatedDataFactory.create_app_mock()

-        # Mock database query to return None
-        mock_query = MagicMock()
-        mock_db_session.session.query.return_value = mock_query
-        mock_query.where.return_value.first.return_value = None
+        # Mock db.session.scalar() to return None
+        mock_db_session.session.scalar.return_value = None

        result = workflow_service.get_draft_workflow(app)

@@ -366,10 +362,8 @@ class TestWorkflowService:
        workflow_id = "workflow-123"
        mock_workflow = TestWorkflowAssociatedDataFactory.create_workflow_mock(version="v1")

-        # Mock database query
-        mock_query = MagicMock()
-        mock_db_session.session.query.return_value = mock_query
-        mock_query.where.return_value.first.return_value = mock_workflow
+        # Mock db.session.scalar() used by get_published_workflow_by_id
+        mock_db_session.session.scalar.return_value = mock_workflow

        result = workflow_service.get_draft_workflow(app, workflow_id=workflow_id)

@@ -384,10 +378,8 @@ class TestWorkflowService:
        workflow_id = "workflow-123"
        mock_workflow = TestWorkflowAssociatedDataFactory.create_workflow_mock(workflow_id=workflow_id, version="v1")

-        # Mock database query
-        mock_query = MagicMock()
-        mock_db_session.session.query.return_value = mock_query
-        mock_query.where.return_value.first.return_value = mock_workflow
+        # Mock db.session.scalar() used by get_published_workflow_by_id
+        mock_db_session.session.scalar.return_value = mock_workflow

        result = workflow_service.get_published_workflow_by_id(app, workflow_id)

@@ -406,10 +398,8 @@ class TestWorkflowService:
            workflow_id=workflow_id, version=Workflow.VERSION_DRAFT
        )

-        # Mock database query
-        mock_query = MagicMock()
-        mock_db_session.session.query.return_value = mock_query
-        mock_query.where.return_value.first.return_value = mock_workflow
+        # Mock db.session.scalar() used by get_published_workflow_by_id
+        mock_db_session.session.scalar.return_value = mock_workflow

        with pytest.raises(IsDraftWorkflowError):
            workflow_service.get_published_workflow_by_id(app, workflow_id)
@@ -419,10 +409,8 @@ class TestWorkflowService:
        app = TestWorkflowAssociatedDataFactory.create_app_mock()
        workflow_id = "nonexistent-workflow"

-        # Mock database query to return None
-        mock_query = MagicMock()
-        mock_db_session.session.query.return_value = mock_query
-        mock_query.where.return_value.first.return_value = None
+        # Mock db.session.scalar() to return None
+        mock_db_session.session.scalar.return_value = None

        result = workflow_service.get_published_workflow_by_id(app, workflow_id)

@@ -434,10 +422,8 @@ class TestWorkflowService:
        app = TestWorkflowAssociatedDataFactory.create_app_mock(workflow_id=workflow_id)
        mock_workflow = TestWorkflowAssociatedDataFactory.create_workflow_mock(workflow_id=workflow_id, version="v1")

-        # Mock database query
-        mock_query = MagicMock()
-        mock_db_session.session.query.return_value = mock_query
-        mock_query.where.return_value.first.return_value = mock_workflow
+        # Mock db.session.scalar() used by get_published_workflow
+        mock_db_session.session.scalar.return_value = mock_workflow

        result = workflow_service.get_published_workflow(app)

@@ -466,11 +452,9 @@ class TestWorkflowService:
        graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph()
        features = {"file_upload": {"enabled": False}}

-        # Mock get_draft_workflow to return None (no existing draft)
+        # Mock db.session.scalar() to return None (no existing draft)
        # This simulates the first time a workflow is created for an app
-        mock_query = MagicMock()
-        mock_db_session.session.query.return_value = mock_query
-        mock_query.where.return_value.first.return_value = None
+        mock_db_session.session.scalar.return_value = None

        with (
            patch.object(workflow_service, "validate_features_structure"),
@@ -504,12 +488,10 @@ class TestWorkflowService:
        features = {"file_upload": {"enabled": False}}
        unique_hash = "test-hash-123"

-        # Mock existing draft workflow
+        # Mock existing draft workflow via db.session.scalar()
        mock_workflow = TestWorkflowAssociatedDataFactory.create_workflow_mock(unique_hash=unique_hash)

-        mock_query = MagicMock()
-        mock_db_session.session.query.return_value = mock_query
-        mock_query.where.return_value.first.return_value = mock_workflow
+        mock_db_session.session.scalar.return_value = mock_workflow

        with (
            patch.object(workflow_service, "validate_features_structure"),
@@ -545,12 +527,10 @@ class TestWorkflowService:
        graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph()
        features = {}

-        # Mock existing draft workflow with different hash
+        # Mock existing draft workflow with different hash via db.session.scalar()
        mock_workflow = TestWorkflowAssociatedDataFactory.create_workflow_mock(unique_hash="old-hash")

-        mock_query = MagicMock()
-        mock_db_session.session.query.return_value = mock_query
-        mock_query.where.return_value.first.return_value = mock_workflow
+        mock_db_session.session.scalar.return_value = mock_workflow

        with pytest.raises(WorkflowHashNotEqualError):
            workflow_service.sync_draft_workflow(
--- a/api/tests/unit_tests/services/tools/test_builtin_tools_manage_service.py
+++ b/api/tests/unit_tests/services/tools/test_builtin_tools_manage_service.py
@@ -347,7 +347,7 @@ class TestGetBuiltinToolProviderCredentials:
    def test_returns_empty_when_no_providers(self, mock_db):
        mock_db.session.no_autoflush.__enter__ = MagicMock(return_value=None)
        mock_db.session.no_autoflush.__exit__ = MagicMock(return_value=False)
-        mock_db.session.query.return_value.filter_by.return_value.order_by.return_value.all.return_value = []
+        mock_db.session.scalars.return_value.all.return_value = []

        result = BuiltinToolManageService.get_builtin_tool_provider_credentials("t", "google")

@@ -362,7 +362,7 @@ class TestGetBuiltinToolProviderCredentials:
        mock_db.session.no_autoflush.__exit__ = MagicMock(return_value=False)

        provider = MagicMock(provider="google", is_default=False)
-        mock_db.session.query.return_value.filter_by.return_value.order_by.return_value.all.return_value = [provider]
+        mock_db.session.scalars.return_value.all.return_value = [provider]

        mock_encrypter = MagicMock()
        mock_encrypter.decrypt.return_value = {"key": "decrypted"}