From 23882a704e59bfaf88956c4163135f489c354f33 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Mon, 20 Jul 2026 19:38:49 +0900 Subject: [PATCH] test: use SQLite sessions in unit misc (#39122) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- .../rag_pipeline/test_rag_pipeline.py | 30 +++++---- ...st_update_provider_when_message_created.py | 61 ++++++++----------- ...ear_free_plan_expired_workflow_run_logs.py | 29 ++++----- 3 files changed, 54 insertions(+), 66 deletions(-) diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py index c39ed13a63c..860e7b6d286 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py @@ -2,6 +2,7 @@ from __future__ import annotations from collections.abc import Iterator from inspect import unwrap +from types import SimpleNamespace from unittest.mock import PropertyMock, patch import pytest @@ -275,9 +276,9 @@ class TestCustomizedPipelineTemplateApi: assert (response, status) == ("", 204) assert deleted_templates == [("template-1", tenant_id)] + @pytest.mark.parametrize("sqlite_session", [(PipelineCustomizedTemplate,)], indirect=True) def test_post_exports_yaml_from_orm_template( - self, - database_app: Flask, + self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine, sqlite_session: Session ) -> None: api = CustomizedPipelineTemplateApi() method = unwrap(api.post) @@ -293,26 +294,29 @@ class TestCustomizedPipelineTemplateApi: language="en-US", created_by="00000000-0000-0000-0000-000000000002", ) - db.session.add(template) - db.session.commit() + template.id = "template-1" + sqlite_session.add(template) + sqlite_session.commit() + monkeypatch.setattr(module, "db", SimpleNamespace(engine=sqlite_engine)) - with database_app.test_request_context("/rag/pipeline/customized/templates/template-1", method="POST"): - response, status = method(api, template.id) + with app.test_request_context("/rag/pipeline/customized/templates/template-1", method="POST"): + response, status = method(api, "template-1") assert status == 200 assert response == {"data": "dsl: value"} + @pytest.mark.parametrize("sqlite_session", [(PipelineCustomizedTemplate,)], indirect=True) def test_post_raises_when_template_is_missing( - self, - database_app: Flask, + self, app: Flask, monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine, sqlite_session: Session ) -> None: api = CustomizedPipelineTemplateApi() method = unwrap(api.post) - with ( - database_app.test_request_context("/rag/pipeline/customized/templates/missing", method="POST"), - pytest.raises(ValueError, match="Customized pipeline template not found"), - ): - method(api, "44444444-4444-4444-4444-444444444444") + assert sqlite_session.get(PipelineCustomizedTemplate, "missing") is None + monkeypatch.setattr(module, "db", SimpleNamespace(engine=sqlite_engine)) + + with app.test_request_context("/rag/pipeline/customized/templates/missing", method="POST"): + with pytest.raises(ValueError, match="Customized pipeline template not found"): + method(api, "missing") class TestPublishCustomizedPipelineTemplateApi: diff --git a/api/tests/unit_tests/events/test_update_provider_when_message_created.py b/api/tests/unit_tests/events/test_update_provider_when_message_created.py index 327c80323b4..a31f0ecdb97 100644 --- a/api/tests/unit_tests/events/test_update_provider_when_message_created.py +++ b/api/tests/unit_tests/events/test_update_provider_when_message_created.py @@ -1,13 +1,12 @@ -from collections.abc import Generator -from contextlib import contextmanager +from collections.abc import Iterator from types import SimpleNamespace from unittest.mock import ANY, patch from uuid import uuid4 import pytest -from sqlalchemy import create_engine, select +from sqlalchemy import select from sqlalchemy.engine import Engine -from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm import Session, sessionmaker from core.app.entities.app_invoke_entities import ChatAppGenerateEntity from core.entities.provider_entities import ProviderQuotaType, QuotaUnit @@ -16,40 +15,29 @@ from models import TenantCreditPool from models.provider import ProviderType -@contextmanager -def _patched_credit_pool_session_factory(engine: Engine) -> Generator[None, None, None]: - session_maker = sessionmaker(bind=engine, expire_on_commit=False) - sessions = [] - - def _session(): - session = session_maker() - sessions.append(session) - return session - - with patch("events.event_handlers.update_provider_when_message_created.db", SimpleNamespace(session=_session)): - try: - yield - finally: - for session in sessions: - session.close() +@pytest.fixture +def credit_pool_session_factory(sqlite_engine: Engine) -> Iterator[sessionmaker[Session]]: + """Bind message-created accounting to fixture-owned SQLite sessions.""" + TenantCreditPool.__table__.create(sqlite_engine) + session_factory = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + with patch("events.event_handlers.update_provider_when_message_created.db.session", session_factory): + yield session_factory -def test_message_created_trial_credit_accounting_does_not_raise_when_balance_is_insufficient() -> None: - engine = create_engine("sqlite:///:memory:") - TenantCreditPool.__table__.create(engine) +def test_message_created_trial_credit_accounting_does_not_raise_when_balance_is_insufficient( + credit_pool_session_factory: sessionmaker[Session], +) -> None: tenant_id = str(uuid4()) pool_id = str(uuid4()) - with engine.begin() as connection: - connection.execute( - TenantCreditPool.__table__.insert(), - { - "id": pool_id, - "tenant_id": tenant_id, - "pool_type": ProviderQuotaType.TRIAL, - "quota_limit": 10, - "quota_used": 9, - }, - ) + pool = TenantCreditPool( + tenant_id=tenant_id, + pool_type=ProviderQuotaType.TRIAL, + quota_limit=10, + quota_used=9, + ) + pool.id = pool_id + with credit_pool_session_factory.begin() as session: + session.add(pool) system_configuration = SimpleNamespace( current_quota_type=ProviderQuotaType.TRIAL, @@ -77,7 +65,6 @@ def test_message_created_trial_credit_accounting_does_not_raise_when_balance_is_ message = SimpleNamespace(message_tokens=2, answer_tokens=1) with ( - _patched_credit_pool_session_factory(engine), patch.object(update_provider_when_message_created, "_execute_provider_updates"), ): update_provider_when_message_created.handle( @@ -85,8 +72,8 @@ def test_message_created_trial_credit_accounting_does_not_raise_when_balance_is_ application_generate_entity=application_generate_entity, ) - with engine.connect() as connection: - quota_used = connection.scalar(select(TenantCreditPool.quota_used).where(TenantCreditPool.id == pool_id)) + with credit_pool_session_factory() as session: + quota_used = session.scalar(select(TenantCreditPool.quota_used).where(TenantCreditPool.id == pool_id)) assert quota_used == 10 diff --git a/api/tests/unit_tests/services/retention/workflow_run/test_clear_free_plan_expired_workflow_run_logs.py b/api/tests/unit_tests/services/retention/workflow_run/test_clear_free_plan_expired_workflow_run_logs.py index 1e15a72f476..524dcd5952f 100644 --- a/api/tests/unit_tests/services/retention/workflow_run/test_clear_free_plan_expired_workflow_run_logs.py +++ b/api/tests/unit_tests/services/retention/workflow_run/test_clear_free_plan_expired_workflow_run_logs.py @@ -6,6 +6,7 @@ import datetime from unittest.mock import MagicMock, patch import pytest +from sqlalchemy.orm import Session from repositories.api_workflow_run_repository import WorkflowRunCleanupRef from services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs import WorkflowRunCleanup @@ -472,24 +473,24 @@ class TestRunDryRunMode: class TestTriggerLogMethods: - def test_delete_trigger_logs(self, cleanup): - session = MagicMock() + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_delete_trigger_logs(self, cleanup, sqlite_session: Session): with patch( "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.SQLAlchemyWorkflowTriggerLogRepository" ) as RepoClass: instance = RepoClass.return_value instance.delete_by_run_ids.return_value = 5 - result = cleanup._delete_trigger_logs(session, ["r1", "r2"]) + result = cleanup._delete_trigger_logs(sqlite_session, ["r1", "r2"]) assert result == 5 - def test_count_trigger_logs(self, cleanup): - session = MagicMock() + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_count_trigger_logs(self, cleanup, sqlite_session: Session): with patch( "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.SQLAlchemyWorkflowTriggerLogRepository" ) as RepoClass: instance = RepoClass.return_value instance.count_by_run_ids.return_value = 3 - result = cleanup._count_trigger_logs(session, ["r1"]) + result = cleanup._count_trigger_logs(sqlite_session, ["r1"]) assert result == 3 @@ -499,26 +500,22 @@ class TestTriggerLogMethods: class TestNodeExecutionMethods: - def test_count_node_executions(self, cleanup): - session = MagicMock() - session.get_bind.return_value = MagicMock() + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_count_node_executions(self, cleanup, sqlite_session: Session): with patch( "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.DifyAPIRepositoryFactory" ) as factory: repo = factory.create_api_workflow_node_execution_repository.return_value repo.count_by_runs.return_value = (10, 2) - with patch("services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.sessionmaker"): - result = cleanup._count_node_executions_by_run_ids(session, ["r1"]) + result = cleanup._count_node_executions_by_run_ids(sqlite_session, ["r1"]) assert result == (10, 2) - def test_delete_node_executions(self, cleanup): - session = MagicMock() - session.get_bind.return_value = MagicMock() + @pytest.mark.parametrize("sqlite_session", [()], indirect=True) + def test_delete_node_executions(self, cleanup, sqlite_session: Session): with patch( "services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.DifyAPIRepositoryFactory" ) as factory: repo = factory.create_api_workflow_node_execution_repository.return_value repo.delete_by_runs.return_value = (5, 1) - with patch("services.retention.workflow_run.clear_free_plan_expired_workflow_run_logs.sessionmaker"): - result = cleanup._delete_node_executions_by_run_ids(session, ["r1"]) + result = cleanup._delete_node_executions_by_run_ids(sqlite_session, ["r1"]) assert result == (5, 1)