refactor(api): use sessionmaker in trigger provider service & dataset… (#34774)

This commit is contained in:
carlos4s
2026-04-08 18:13:38 -05:00
committed by GitHub
parent 3a4756449a
commit 4c70bfa8b8
4 changed files with 22 additions and 27 deletions

View File

@ -4909,15 +4909,17 @@ class TestInternalHooksCoverage:
session_ctx.__enter__.return_value = session
session_ctx.__exit__.return_value = False
sessionmaker_ctx = MagicMock()
sessionmaker_ctx.begin.return_value = session_ctx
with (
patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=Mock())),
patch("core.rag.retrieval.dataset_retrieval.Session", return_value=session_ctx),
patch("core.rag.retrieval.dataset_retrieval.sessionmaker", return_value=sessionmaker_ctx),
patch.object(retrieval, "_send_trace_task") as mock_trace,
):
retrieval._on_retrieval_end(flask_app=app, documents=docs, message_id="m1", timer={"cost": 1})
query.update.assert_called_once()
session.commit.assert_called_once()
mock_trace.assert_called_once()
def test_retriever_variants(self, retrieval: DatasetRetrieval) -> None:

View File

@ -63,6 +63,12 @@ def mock_session(mocker: MockerFixture) -> MagicMock:
mock_session_cm.__enter__.return_value = mock_session_instance
mock_session_cm.__exit__.return_value = False
mocker.patch("services.trigger.trigger_provider_service.Session", return_value=mock_session_cm)
mock_begin_cm = MagicMock()
mock_begin_cm.__enter__.return_value = mock_session_instance
mock_begin_cm.__exit__.return_value = False
mock_sessionmaker_instance = MagicMock()
mock_sessionmaker_instance.begin.return_value = mock_begin_cm
mocker.patch("services.trigger.trigger_provider_service.sessionmaker", return_value=mock_sessionmaker_instance)
return mock_session_instance
@ -212,7 +218,6 @@ def test_add_trigger_subscription_should_create_subscription_successfully_for_ap
# Assert
assert result["result"] == "success"
mock_session.add.assert_called_once()
mock_session.commit.assert_called_once()
def test_add_trigger_subscription_should_store_empty_credentials_for_unauthorized_type(
@ -406,7 +411,7 @@ def test_update_trigger_subscription_should_update_fields_and_clear_cache(
assert subscription.credentials == {"api_key": "new-key"}
assert subscription.credential_expires_at == 100
assert subscription.expires_at == 200
mock_session.commit.assert_called_once()
mock_delete_cache.assert_called_once()
@ -593,7 +598,7 @@ def test_refresh_oauth_token_should_refresh_and_persist_new_credentials(
assert result == {"result": "success", "expires_at": 12345}
assert subscription.credentials == {"access_token": "new"}
assert subscription.credential_expires_at == 12345
mock_session.commit.assert_called_once()
cache.delete.assert_called_once()
@ -664,7 +669,7 @@ def test_refresh_subscription_should_refresh_and_persist_properties(
assert result == {"result": "success", "expires_at": 999}
assert subscription.properties == {"p": "new-enc"}
assert subscription.expires_at == 999
mock_session.commit.assert_called_once()
prop_cache.delete.assert_called_once()
@ -838,7 +843,6 @@ def test_save_custom_oauth_client_params_should_create_record_and_clear_params_w
assert fake_model.encrypted_oauth_params == "{}"
assert fake_model.enabled is True
mock_session.add.assert_called_once_with(fake_model)
mock_session.commit.assert_called_once()
def test_save_custom_oauth_client_params_should_merge_hidden_values_and_delete_cache(
@ -870,7 +874,6 @@ def test_save_custom_oauth_client_params_should_merge_hidden_values_and_delete_c
assert result == {"result": "success"}
assert json.loads(custom_client.encrypted_oauth_params) == {"client_id": "new-id"}
cache.delete.assert_called_once()
mock_session.commit.assert_called_once()
def test_get_custom_oauth_client_params_should_return_empty_when_record_missing(
@ -921,7 +924,6 @@ def test_delete_custom_oauth_client_params_should_delete_record_and_commit(
# Assert
assert result == {"result": "success"}
mock_session.commit.assert_called_once()
@pytest.mark.parametrize("exists", [True, False])