opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-21 10:38:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

test: use SQLite sessions in unit misc (#39122)

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Asuka Minato 2026-07-20 19:38:49 +09:00 committed by GitHub
parent 1bd5254ea7
commit 23882a704e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 54 additions and 66 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline.py
View File

@ -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:

61
api/tests/unit_tests/events/test_update_provider_when_message_created.py
View File

@ -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

29
api/tests/unit_tests/services/retention/workflow_run/test_clear_free_plan_expired_workflow_run_logs.py
View File

@ -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)