From 608958de1c43eba96b1e88a6b046c69a6a2e2327 Mon Sep 17 00:00:00 2001 From: Renzo <170978465+RenzoMXD@users.noreply.github.com> Date: Fri, 3 Apr 2026 05:42:16 +0200 Subject: [PATCH] refactor: select in external_knowledge_service (#34493) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- api/services/external_knowledge_service.py | 70 +++-- .../services/external_dataset_service.py | 46 +-- .../services/test_external_dataset_service.py | 292 ++---------------- 3 files changed, 101 insertions(+), 307 deletions(-) diff --git a/api/services/external_knowledge_service.py b/api/services/external_knowledge_service.py index 9a522ece528..b178718baae 100644 --- a/api/services/external_knowledge_service.py +++ b/api/services/external_knowledge_service.py @@ -5,7 +5,7 @@ from urllib.parse import urlparse import httpx from graphon.nodes.http_request.exc import InvalidHttpMethodError -from sqlalchemy import select +from sqlalchemy import func, select from constants import HIDDEN_VALUE from core.helper import ssrf_proxy @@ -103,8 +103,10 @@ class ExternalDatasetService: @staticmethod def get_external_knowledge_api(external_knowledge_api_id: str, tenant_id: str) -> ExternalKnowledgeApis: - external_knowledge_api: ExternalKnowledgeApis | None = ( - db.session.query(ExternalKnowledgeApis).filter_by(id=external_knowledge_api_id, tenant_id=tenant_id).first() + external_knowledge_api: ExternalKnowledgeApis | None = db.session.scalar( + select(ExternalKnowledgeApis) + .where(ExternalKnowledgeApis.id == external_knowledge_api_id, ExternalKnowledgeApis.tenant_id == tenant_id) + .limit(1) ) if external_knowledge_api is None: raise ValueError("api template not found") @@ -112,8 +114,10 @@ class ExternalDatasetService: @staticmethod def update_external_knowledge_api(tenant_id, user_id, external_knowledge_api_id, args) -> ExternalKnowledgeApis: - external_knowledge_api: ExternalKnowledgeApis | None = ( - db.session.query(ExternalKnowledgeApis).filter_by(id=external_knowledge_api_id, tenant_id=tenant_id).first() + external_knowledge_api: ExternalKnowledgeApis | None = db.session.scalar( + select(ExternalKnowledgeApis) + .where(ExternalKnowledgeApis.id == external_knowledge_api_id, ExternalKnowledgeApis.tenant_id == tenant_id) + .limit(1) ) if external_knowledge_api is None: raise ValueError("api template not found") @@ -132,8 +136,10 @@ class ExternalDatasetService: @staticmethod def delete_external_knowledge_api(tenant_id: str, external_knowledge_api_id: str): - external_knowledge_api = ( - db.session.query(ExternalKnowledgeApis).filter_by(id=external_knowledge_api_id, tenant_id=tenant_id).first() + external_knowledge_api = db.session.scalar( + select(ExternalKnowledgeApis) + .where(ExternalKnowledgeApis.id == external_knowledge_api_id, ExternalKnowledgeApis.tenant_id == tenant_id) + .limit(1) ) if external_knowledge_api is None: raise ValueError("api template not found") @@ -144,9 +150,12 @@ class ExternalDatasetService: @staticmethod def external_knowledge_api_use_check(external_knowledge_api_id: str) -> tuple[bool, int]: count = ( - db.session.query(ExternalKnowledgeBindings) - .filter_by(external_knowledge_api_id=external_knowledge_api_id) - .count() + db.session.scalar( + select(func.count(ExternalKnowledgeBindings.id)).where( + ExternalKnowledgeBindings.external_knowledge_api_id == external_knowledge_api_id + ) + ) + or 0 ) if count > 0: return True, count @@ -154,8 +163,10 @@ class ExternalDatasetService: @staticmethod def get_external_knowledge_binding_with_dataset_id(tenant_id: str, dataset_id: str) -> ExternalKnowledgeBindings: - external_knowledge_binding: ExternalKnowledgeBindings | None = ( - db.session.query(ExternalKnowledgeBindings).filter_by(dataset_id=dataset_id, tenant_id=tenant_id).first() + external_knowledge_binding: ExternalKnowledgeBindings | None = db.session.scalar( + select(ExternalKnowledgeBindings) + .where(ExternalKnowledgeBindings.dataset_id == dataset_id, ExternalKnowledgeBindings.tenant_id == tenant_id) + .limit(1) ) if not external_knowledge_binding: raise ValueError("external knowledge binding not found") @@ -163,8 +174,10 @@ class ExternalDatasetService: @staticmethod def document_create_args_validate(tenant_id: str, external_knowledge_api_id: str, process_parameter: dict): - external_knowledge_api = ( - db.session.query(ExternalKnowledgeApis).filter_by(id=external_knowledge_api_id, tenant_id=tenant_id).first() + external_knowledge_api = db.session.scalar( + select(ExternalKnowledgeApis) + .where(ExternalKnowledgeApis.id == external_knowledge_api_id, ExternalKnowledgeApis.tenant_id == tenant_id) + .limit(1) ) if external_knowledge_api is None or external_knowledge_api.settings is None: raise ValueError("api template not found") @@ -238,12 +251,17 @@ class ExternalDatasetService: @staticmethod def create_external_dataset(tenant_id: str, user_id: str, args: dict) -> Dataset: # check if dataset name already exists - if db.session.query(Dataset).filter_by(name=args.get("name"), tenant_id=tenant_id).first(): + if db.session.scalar( + select(Dataset).where(Dataset.name == args.get("name"), Dataset.tenant_id == tenant_id).limit(1) + ): raise DatasetNameDuplicateError(f"Dataset with name {args.get('name')} already exists.") - external_knowledge_api = ( - db.session.query(ExternalKnowledgeApis) - .filter_by(id=args.get("external_knowledge_api_id"), tenant_id=tenant_id) - .first() + external_knowledge_api = db.session.scalar( + select(ExternalKnowledgeApis) + .where( + ExternalKnowledgeApis.id == args.get("external_knowledge_api_id"), + ExternalKnowledgeApis.tenant_id == tenant_id, + ) + .limit(1) ) if external_knowledge_api is None: @@ -286,16 +304,18 @@ class ExternalDatasetService: external_retrieval_parameters: dict, metadata_condition: MetadataCondition | None = None, ): - external_knowledge_binding = ( - db.session.query(ExternalKnowledgeBindings).filter_by(dataset_id=dataset_id, tenant_id=tenant_id).first() + external_knowledge_binding = db.session.scalar( + select(ExternalKnowledgeBindings) + .where(ExternalKnowledgeBindings.dataset_id == dataset_id, ExternalKnowledgeBindings.tenant_id == tenant_id) + .limit(1) ) if not external_knowledge_binding: raise ValueError("external knowledge binding not found") - external_knowledge_api = ( - db.session.query(ExternalKnowledgeApis) - .filter_by(id=external_knowledge_binding.external_knowledge_api_id) - .first() + external_knowledge_api = db.session.scalar( + select(ExternalKnowledgeApis) + .where(ExternalKnowledgeApis.id == external_knowledge_binding.external_knowledge_api_id) + .limit(1) ) if external_knowledge_api is None or external_knowledge_api.settings is None: raise ValueError("external api template not found") diff --git a/api/tests/unit_tests/services/external_dataset_service.py b/api/tests/unit_tests/services/external_dataset_service.py index 70bd1c73b30..5848603ab8e 100644 --- a/api/tests/unit_tests/services/external_dataset_service.py +++ b/api/tests/unit_tests/services/external_dataset_service.py @@ -292,7 +292,7 @@ class TestExternalDatasetServiceCrudExternalKnowledgeApi: """ api = Mock(spec=ExternalKnowledgeApis) - mock_db_session.query.return_value.filter_by.return_value.first.return_value = api + mock_db_session.scalar.return_value = api result = ExternalDatasetService.get_external_knowledge_api("api-id", "tenant-id") assert result is api @@ -302,7 +302,7 @@ class TestExternalDatasetServiceCrudExternalKnowledgeApi: When the record is absent, a ``ValueError`` is raised. """ - mock_db_session.query.return_value.filter_by.return_value.first.return_value = None + mock_db_session.scalar.return_value = None with pytest.raises(ValueError, match="api template not found"): ExternalDatasetService.get_external_knowledge_api("missing-id", "tenant-id") @@ -320,7 +320,7 @@ class TestExternalDatasetServiceCrudExternalKnowledgeApi: existing_api = Mock(spec=ExternalKnowledgeApis) existing_api.settings_dict = {"api_key": "stored-key"} existing_api.settings = '{"api_key":"stored-key"}' - mock_db_session.query.return_value.filter_by.return_value.first.return_value = existing_api + mock_db_session.scalar.return_value = existing_api args = { "name": "New Name", @@ -340,7 +340,7 @@ class TestExternalDatasetServiceCrudExternalKnowledgeApi: Updating a non‑existent API template should raise ``ValueError``. """ - mock_db_session.query.return_value.filter_by.return_value.first.return_value = None + mock_db_session.scalar.return_value = None with pytest.raises(ValueError, match="api template not found"): ExternalDatasetService.update_external_knowledge_api( @@ -356,7 +356,7 @@ class TestExternalDatasetServiceCrudExternalKnowledgeApi: """ api = Mock(spec=ExternalKnowledgeApis) - mock_db_session.query.return_value.filter_by.return_value.first.return_value = api + mock_db_session.scalar.return_value = api ExternalDatasetService.delete_external_knowledge_api("tenant-1", "api-1") @@ -368,7 +368,7 @@ class TestExternalDatasetServiceCrudExternalKnowledgeApi: Deletion of a missing template should raise ``ValueError``. """ - mock_db_session.query.return_value.filter_by.return_value.first.return_value = None + mock_db_session.scalar.return_value = None with pytest.raises(ValueError, match="api template not found"): ExternalDatasetService.delete_external_knowledge_api("tenant-1", "missing") @@ -394,7 +394,7 @@ class TestExternalDatasetServiceUsageAndBindings: When there are bindings, ``external_knowledge_api_use_check`` returns True and count. """ - mock_db_session.query.return_value.filter_by.return_value.count.return_value = 3 + mock_db_session.scalar.return_value = 3 in_use, count = ExternalDatasetService.external_knowledge_api_use_check("api-1") @@ -406,7 +406,7 @@ class TestExternalDatasetServiceUsageAndBindings: Zero bindings should return ``(False, 0)``. """ - mock_db_session.query.return_value.filter_by.return_value.count.return_value = 0 + mock_db_session.scalar.return_value = 0 in_use, count = ExternalDatasetService.external_knowledge_api_use_check("api-1") @@ -419,7 +419,7 @@ class TestExternalDatasetServiceUsageAndBindings: """ binding = Mock(spec=ExternalKnowledgeBindings) - mock_db_session.query.return_value.filter_by.return_value.first.return_value = binding + mock_db_session.scalar.return_value = binding result = ExternalDatasetService.get_external_knowledge_binding_with_dataset_id("tenant-1", "ds-1") assert result is binding @@ -429,7 +429,7 @@ class TestExternalDatasetServiceUsageAndBindings: Missing binding should result in a ``ValueError``. """ - mock_db_session.query.return_value.filter_by.return_value.first.return_value = None + mock_db_session.scalar.return_value = None with pytest.raises(ValueError, match="external knowledge binding not found"): ExternalDatasetService.get_external_knowledge_binding_with_dataset_id("tenant-1", "ds-1") @@ -460,7 +460,7 @@ class TestExternalDatasetServiceDocumentCreateArgsValidate: '[{"document_process_setting":[{"name":"foo","required":true},{"name":"bar","required":false}]}]' ) # Raw string; the service itself calls json.loads on it - mock_db_session.query.return_value.filter_by.return_value.first.return_value = external_api + mock_db_session.scalar.return_value = external_api process_parameter = {"foo": "value", "bar": "optional"} @@ -474,7 +474,7 @@ class TestExternalDatasetServiceDocumentCreateArgsValidate: When the referenced API template is missing, a ``ValueError`` is raised. """ - mock_db_session.query.return_value.filter_by.return_value.first.return_value = None + mock_db_session.scalar.return_value = None with pytest.raises(ValueError, match="api template not found"): ExternalDatasetService.document_create_args_validate("tenant-1", "missing", {}) @@ -488,7 +488,7 @@ class TestExternalDatasetServiceDocumentCreateArgsValidate: external_api.settings = ( '[{"document_process_setting":[{"name":"foo","required":true},{"name":"bar","required":false}]}]' ) - mock_db_session.query.return_value.filter_by.return_value.first.return_value = external_api + mock_db_session.scalar.return_value = external_api process_parameter = {"bar": "present"} # missing "foo" @@ -702,7 +702,7 @@ class TestExternalDatasetServiceCreateExternalDataset: } # No existing dataset with same name. - mock_db_session.query.return_value.filter_by.return_value.first.side_effect = [ + mock_db_session.scalar.side_effect = [ None, # duplicate‑name check Mock(spec=ExternalKnowledgeApis), # external knowledge api ] @@ -724,7 +724,7 @@ class TestExternalDatasetServiceCreateExternalDataset: """ existing_dataset = Mock(spec=Dataset) - mock_db_session.query.return_value.filter_by.return_value.first.return_value = existing_dataset + mock_db_session.scalar.return_value = existing_dataset args = { "name": "Existing", @@ -744,7 +744,7 @@ class TestExternalDatasetServiceCreateExternalDataset: """ # First call: duplicate name check – not found. - mock_db_session.query.return_value.filter_by.return_value.first.side_effect = [ + mock_db_session.scalar.side_effect = [ None, None, # external knowledge api lookup ] @@ -763,8 +763,10 @@ class TestExternalDatasetServiceCreateExternalDataset: ``external_knowledge_id`` and ``external_knowledge_api_id`` are mandatory. """ - # duplicate name check - mock_db_session.query.return_value.filter_by.return_value.first.side_effect = [ + # duplicate name check — two calls to create_external_dataset, each does 2 scalar calls + mock_db_session.scalar.side_effect = [ + None, + Mock(spec=ExternalKnowledgeApis), None, Mock(spec=ExternalKnowledgeApis), ] @@ -826,7 +828,7 @@ class TestExternalDatasetServiceFetchExternalKnowledgeRetrieval: api.settings = '{"endpoint":"https://example.com","api_key":"secret"}' # First query: binding; second query: api. - mock_db_session.query.return_value.filter_by.return_value.first.side_effect = [ + mock_db_session.scalar.side_effect = [ binding, api, ] @@ -861,7 +863,7 @@ class TestExternalDatasetServiceFetchExternalKnowledgeRetrieval: Missing binding should raise ``ValueError``. """ - mock_db_session.query.return_value.filter_by.return_value.first.return_value = None + mock_db_session.scalar.return_value = None with pytest.raises(ValueError, match="external knowledge binding not found"): ExternalDatasetService.fetch_external_knowledge_retrieval( @@ -878,7 +880,7 @@ class TestExternalDatasetServiceFetchExternalKnowledgeRetrieval: """ binding = ExternalDatasetTestDataFactory.create_external_binding() - mock_db_session.query.return_value.filter_by.return_value.first.side_effect = [ + mock_db_session.scalar.side_effect = [ binding, None, ] @@ -901,7 +903,7 @@ class TestExternalDatasetServiceFetchExternalKnowledgeRetrieval: api = Mock(spec=ExternalKnowledgeApis) api.settings = '{"endpoint":"https://example.com","api_key":"secret"}' - mock_db_session.query.return_value.filter_by.return_value.first.side_effect = [ + mock_db_session.scalar.side_effect = [ binding, api, ] diff --git a/api/tests/unit_tests/services/test_external_dataset_service.py b/api/tests/unit_tests/services/test_external_dataset_service.py index 3709e1fa94d..7c8dab5029e 100644 --- a/api/tests/unit_tests/services/test_external_dataset_service.py +++ b/api/tests/unit_tests/services/test_external_dataset_service.py @@ -799,10 +799,7 @@ class TestExternalDatasetServiceGetAPI: api_id = "api-123" expected_api = factory.create_external_knowledge_api_mock(api_id=api_id) - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = expected_api + mock_db.session.scalar.return_value = expected_api # Act tenant_id = "tenant-123" @@ -810,16 +807,12 @@ class TestExternalDatasetServiceGetAPI: # Assert assert result.id == api_id - mock_query.filter_by.assert_called_once_with(id=api_id, tenant_id=tenant_id) @patch("services.external_knowledge_service.db") def test_get_external_knowledge_api_not_found(self, mock_db, factory): """Test error when API is not found.""" # Arrange - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = None + mock_db.session.scalar.return_value = None # Act & Assert with pytest.raises(ValueError, match="api template not found"): @@ -848,10 +841,7 @@ class TestExternalDatasetServiceUpdateAPI: "settings": {"endpoint": "https://new.example.com", "api_key": "new-key"}, } - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = existing_api + mock_db.session.scalar.return_value = existing_api # Act result = ExternalDatasetService.update_external_knowledge_api(tenant_id, user_id, api_id, args) @@ -881,10 +871,7 @@ class TestExternalDatasetServiceUpdateAPI: "settings": {"endpoint": "https://api.example.com", "api_key": HIDDEN_VALUE}, } - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = existing_api + mock_db.session.scalar.return_value = existing_api # Act result = ExternalDatasetService.update_external_knowledge_api(tenant_id, "user-123", api_id, args) @@ -897,10 +884,7 @@ class TestExternalDatasetServiceUpdateAPI: def test_update_external_knowledge_api_not_found(self, mock_db, factory): """Test error when API is not found.""" # Arrange - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = None + mock_db.session.scalar.return_value = None args = {"name": "Updated API"} @@ -912,10 +896,7 @@ class TestExternalDatasetServiceUpdateAPI: def test_update_external_knowledge_api_tenant_mismatch(self, mock_db, factory): """Test error when tenant ID doesn't match.""" # Arrange - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = None + mock_db.session.scalar.return_value = None args = {"name": "Updated API"} @@ -934,10 +915,7 @@ class TestExternalDatasetServiceUpdateAPI: args = {"name": "New Name Only"} - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = existing_api + mock_db.session.scalar.return_value = existing_api # Act result = ExternalDatasetService.update_external_knowledge_api("tenant-123", "user-123", "api-123", args) @@ -958,10 +936,7 @@ class TestExternalDatasetServiceDeleteAPI: existing_api = factory.create_external_knowledge_api_mock(api_id=api_id, tenant_id=tenant_id) - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = existing_api + mock_db.session.scalar.return_value = existing_api # Act ExternalDatasetService.delete_external_knowledge_api(tenant_id, api_id) @@ -974,10 +949,7 @@ class TestExternalDatasetServiceDeleteAPI: def test_delete_external_knowledge_api_not_found(self, mock_db, factory): """Test error when API is not found.""" # Arrange - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = None + mock_db.session.scalar.return_value = None # Act & Assert with pytest.raises(ValueError, match="api template not found"): @@ -987,10 +959,7 @@ class TestExternalDatasetServiceDeleteAPI: def test_delete_external_knowledge_api_tenant_mismatch(self, mock_db, factory): """Test error when tenant ID doesn't match.""" # Arrange - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = None + mock_db.session.scalar.return_value = None # Act & Assert with pytest.raises(ValueError, match="api template not found"): @@ -1006,10 +975,7 @@ class TestExternalDatasetServiceAPIUseCheck: # Arrange api_id = "api-123" - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.count.return_value = 1 + mock_db.session.scalar.return_value = 1 # Act in_use, count = ExternalDatasetService.external_knowledge_api_use_check(api_id) @@ -1024,10 +990,7 @@ class TestExternalDatasetServiceAPIUseCheck: # Arrange api_id = "api-123" - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.count.return_value = 10 + mock_db.session.scalar.return_value = 10 # Act in_use, count = ExternalDatasetService.external_knowledge_api_use_check(api_id) @@ -1042,10 +1005,7 @@ class TestExternalDatasetServiceAPIUseCheck: # Arrange api_id = "api-123" - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.count.return_value = 0 + mock_db.session.scalar.return_value = 0 # Act in_use, count = ExternalDatasetService.external_knowledge_api_use_check(api_id) @@ -1067,10 +1027,7 @@ class TestExternalDatasetServiceGetBinding: expected_binding = factory.create_external_knowledge_binding_mock(tenant_id=tenant_id, dataset_id=dataset_id) - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = expected_binding + mock_db.session.scalar.return_value = expected_binding # Act result = ExternalDatasetService.get_external_knowledge_binding_with_dataset_id(tenant_id, dataset_id) @@ -1083,10 +1040,7 @@ class TestExternalDatasetServiceGetBinding: def test_get_external_knowledge_binding_not_found(self, mock_db, factory): """Test error when binding is not found.""" # Arrange - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = None + mock_db.session.scalar.return_value = None # Act & Assert with pytest.raises(ValueError, match="external knowledge binding not found"): @@ -1113,10 +1067,7 @@ class TestExternalDatasetServiceDocumentValidate: api = factory.create_external_knowledge_api_mock(api_id=api_id, settings=[settings]) - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = api + mock_db.session.scalar.return_value = api process_parameter = {"param1": "value1", "param2": "value2"} @@ -1134,10 +1085,7 @@ class TestExternalDatasetServiceDocumentValidate: api = factory.create_external_knowledge_api_mock(api_id=api_id, settings=[settings]) - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = api + mock_db.session.scalar.return_value = api process_parameter = {} @@ -1149,10 +1097,7 @@ class TestExternalDatasetServiceDocumentValidate: def test_document_create_args_validate_api_not_found(self, mock_db, factory): """Test validation fails when API is not found.""" # Arrange - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = None + mock_db.session.scalar.return_value = None # Act & Assert with pytest.raises(ValueError, match="api template not found"): @@ -1165,10 +1110,7 @@ class TestExternalDatasetServiceDocumentValidate: settings = {} api = factory.create_external_knowledge_api_mock(settings=[settings]) - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = api + mock_db.session.scalar.return_value = api # Act & Assert - should not raise ExternalDatasetService.document_create_args_validate("tenant-123", "api-123", {}) @@ -1186,10 +1128,7 @@ class TestExternalDatasetServiceDocumentValidate: api = factory.create_external_knowledge_api_mock(settings=[settings]) - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = api + mock_db.session.scalar.return_value = api process_parameter = {"required_param": "value"} @@ -1498,24 +1437,7 @@ class TestExternalDatasetServiceCreateDataset: api = factory.create_external_knowledge_api_mock(api_id="api-123") - # Mock database queries - mock_dataset_query = MagicMock() - mock_api_query = MagicMock() - - def query_side_effect(model): - if model == Dataset: - return mock_dataset_query - elif model == ExternalKnowledgeApis: - return mock_api_query - return MagicMock() - - mock_db.session.query.side_effect = query_side_effect - - mock_dataset_query.filter_by.return_value = mock_dataset_query - mock_dataset_query.first.return_value = None - - mock_api_query.filter_by.return_value = mock_api_query - mock_api_query.first.return_value = api + mock_db.session.scalar.side_effect = [None, api] # Act result = ExternalDatasetService.create_external_dataset(tenant_id, user_id, args) @@ -1534,10 +1456,7 @@ class TestExternalDatasetServiceCreateDataset: # Arrange existing_dataset = factory.create_dataset_mock(name="Duplicate Dataset") - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = existing_dataset + mock_db.session.scalar.return_value = existing_dataset args = {"name": "Duplicate Dataset"} @@ -1549,23 +1468,7 @@ class TestExternalDatasetServiceCreateDataset: def test_create_external_dataset_api_not_found_error(self, mock_db, factory): """Test error when external knowledge API is not found.""" # Arrange - mock_dataset_query = MagicMock() - mock_api_query = MagicMock() - - def query_side_effect(model): - if model == Dataset: - return mock_dataset_query - elif model == ExternalKnowledgeApis: - return mock_api_query - return MagicMock() - - mock_db.session.query.side_effect = query_side_effect - - mock_dataset_query.filter_by.return_value = mock_dataset_query - mock_dataset_query.first.return_value = None - - mock_api_query.filter_by.return_value = mock_api_query - mock_api_query.first.return_value = None + mock_db.session.scalar.side_effect = [None, None] args = {"name": "Test Dataset", "external_knowledge_api_id": "nonexistent-api"} @@ -1579,23 +1482,7 @@ class TestExternalDatasetServiceCreateDataset: # Arrange api = factory.create_external_knowledge_api_mock() - mock_dataset_query = MagicMock() - mock_api_query = MagicMock() - - def query_side_effect(model): - if model == Dataset: - return mock_dataset_query - elif model == ExternalKnowledgeApis: - return mock_api_query - return MagicMock() - - mock_db.session.query.side_effect = query_side_effect - - mock_dataset_query.filter_by.return_value = mock_dataset_query - mock_dataset_query.first.return_value = None - - mock_api_query.filter_by.return_value = mock_api_query - mock_api_query.first.return_value = api + mock_db.session.scalar.side_effect = [None, api] args = {"name": "Test Dataset", "external_knowledge_api_id": "api-123"} @@ -1609,23 +1496,7 @@ class TestExternalDatasetServiceCreateDataset: # Arrange api = factory.create_external_knowledge_api_mock() - mock_dataset_query = MagicMock() - mock_api_query = MagicMock() - - def query_side_effect(model): - if model == Dataset: - return mock_dataset_query - elif model == ExternalKnowledgeApis: - return mock_api_query - return MagicMock() - - mock_db.session.query.side_effect = query_side_effect - - mock_dataset_query.filter_by.return_value = mock_dataset_query - mock_dataset_query.first.return_value = None - - mock_api_query.filter_by.return_value = mock_api_query - mock_api_query.first.return_value = api + mock_db.session.scalar.side_effect = [None, api] args = {"name": "Test Dataset", "external_knowledge_id": "knowledge-123"} @@ -1651,23 +1522,7 @@ class TestExternalDatasetServiceFetchRetrieval: ) api = factory.create_external_knowledge_api_mock(api_id="api-123") - mock_binding_query = MagicMock() - mock_api_query = MagicMock() - - def query_side_effect(model): - if model == ExternalKnowledgeBindings: - return mock_binding_query - elif model == ExternalKnowledgeApis: - return mock_api_query - return MagicMock() - - mock_db.session.query.side_effect = query_side_effect - - mock_binding_query.filter_by.return_value = mock_binding_query - mock_binding_query.first.return_value = binding - - mock_api_query.filter_by.return_value = mock_api_query - mock_api_query.first.return_value = api + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = 200 @@ -1695,10 +1550,7 @@ class TestExternalDatasetServiceFetchRetrieval: def test_fetch_external_knowledge_retrieval_binding_not_found_error(self, mock_db, factory): """Test error when external knowledge binding is not found.""" # Arrange - mock_query = MagicMock() - mock_db.session.query.return_value = mock_query - mock_query.filter_by.return_value = mock_query - mock_query.first.return_value = None + mock_db.session.scalar.return_value = None # Act & Assert with pytest.raises(ValueError, match="external knowledge binding not found"): @@ -1712,23 +1564,7 @@ class TestExternalDatasetServiceFetchRetrieval: binding = factory.create_external_knowledge_binding_mock() api = factory.create_external_knowledge_api_mock() - mock_binding_query = MagicMock() - mock_api_query = MagicMock() - - def query_side_effect(model): - if model == ExternalKnowledgeBindings: - return mock_binding_query - elif model == ExternalKnowledgeApis: - return mock_api_query - return MagicMock() - - mock_db.session.query.side_effect = query_side_effect - - mock_binding_query.filter_by.return_value = mock_binding_query - mock_binding_query.first.return_value = binding - - mock_api_query.filter_by.return_value = mock_api_query - mock_api_query.first.return_value = api + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = 200 @@ -1751,23 +1587,7 @@ class TestExternalDatasetServiceFetchRetrieval: binding = factory.create_external_knowledge_binding_mock() api = factory.create_external_knowledge_api_mock() - mock_binding_query = MagicMock() - mock_api_query = MagicMock() - - def query_side_effect(model): - if model == ExternalKnowledgeBindings: - return mock_binding_query - elif model == ExternalKnowledgeApis: - return mock_api_query - return MagicMock() - - mock_db.session.query.side_effect = query_side_effect - - mock_binding_query.filter_by.return_value = mock_binding_query - mock_binding_query.first.return_value = binding - - mock_api_query.filter_by.return_value = mock_api_query - mock_api_query.first.return_value = api + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = 200 @@ -1799,23 +1619,7 @@ class TestExternalDatasetServiceFetchRetrieval: binding = factory.create_external_knowledge_binding_mock() api = factory.create_external_knowledge_api_mock() - mock_binding_query = MagicMock() - mock_api_query = MagicMock() - - def query_side_effect(model): - if model == ExternalKnowledgeBindings: - return mock_binding_query - elif model == ExternalKnowledgeApis: - return mock_api_query - return MagicMock() - - mock_db.session.query.side_effect = query_side_effect - - mock_binding_query.filter_by.return_value = mock_binding_query - mock_binding_query.first.return_value = binding - - mock_api_query.filter_by.return_value = mock_api_query - mock_api_query.first.return_value = api + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = 500 @@ -1856,23 +1660,7 @@ class TestExternalDatasetServiceFetchRetrieval: ) api = factory.create_external_knowledge_api_mock(api_id="api-123") - mock_binding_query = MagicMock() - mock_api_query = MagicMock() - - def query_side_effect(model): - if model == ExternalKnowledgeBindings: - return mock_binding_query - elif model == ExternalKnowledgeApis: - return mock_api_query - return MagicMock() - - mock_db.session.query.side_effect = query_side_effect - - mock_binding_query.filter_by.return_value = mock_binding_query - mock_binding_query.first.return_value = binding - - mock_api_query.filter_by.return_value = mock_api_query - mock_api_query.first.return_value = api + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = status_code @@ -1891,23 +1679,7 @@ class TestExternalDatasetServiceFetchRetrieval: binding = factory.create_external_knowledge_binding_mock() api = factory.create_external_knowledge_api_mock() - mock_binding_query = MagicMock() - mock_api_query = MagicMock() - - def query_side_effect(model): - if model == ExternalKnowledgeBindings: - return mock_binding_query - elif model == ExternalKnowledgeApis: - return mock_api_query - return MagicMock() - - mock_db.session.query.side_effect = query_side_effect - - mock_binding_query.filter_by.return_value = mock_binding_query - mock_binding_query.first.return_value = binding - - mock_api_query.filter_by.return_value = mock_api_query - mock_api_query.first.return_value = api + mock_db.session.scalar.side_effect = [binding, api] mock_response = MagicMock() mock_response.status_code = 503