mirror of
https://github.com/langgenius/dify.git
synced 2026-09-24 23:22:26 +08:00
test: use SQLite sessions in commands (#39082)
This commit is contained in:
Asuka Minato
@@ -1,12 +1,35 @@
|
||||
"""Tests for workflow-run archive command database boundaries.
|
||||
|
||||
Planning deliberately creates a fresh session for every tenant prefix and for
|
||||
every database retry. SQLite-backed tests keep the query, filtering, counting,
|
||||
and session lifecycle real while clocks, billing lookup, and command failures
|
||||
remain narrow external-boundary substitutions.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
from unittest.mock import MagicMock
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, Mock
|
||||
from uuid import uuid4
|
||||
|
||||
import click
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
from sqlalchemy import Engine, event, text
|
||||
from sqlalchemy.exc import OperationalError
|
||||
from sqlalchemy.orm import Session, scoped_session, sessionmaker
|
||||
|
||||
from commands import retention
|
||||
from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus
|
||||
from models.base import TypeBase
|
||||
from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom
|
||||
from models.workflow import (
|
||||
WorkflowNodeExecutionModel,
|
||||
WorkflowNodeExecutionTriggeredFrom,
|
||||
WorkflowRun,
|
||||
WorkflowType,
|
||||
)
|
||||
from services.retention.workflow_run import bundle_archive_maintenance
|
||||
from services.retention.workflow_run.bundle_archive_maintenance import (
|
||||
BundleOperationResult,
|
||||
@@ -27,11 +50,74 @@ def _db_disconnect_error() -> OperationalError:
|
||||
)
|
||||
|
||||
|
||||
def _session_context(session):
|
||||
context = MagicMock()
|
||||
context.__enter__.return_value = session
|
||||
context.__exit__.return_value = False
|
||||
return context
|
||||
@dataclass(frozen=True)
|
||||
class ArchiveDatabase:
|
||||
"""Creates archive candidates and their node executions in SQLite."""
|
||||
|
||||
session_maker: sessionmaker[Session]
|
||||
end_before: datetime.datetime
|
||||
|
||||
def add_run(
|
||||
self,
|
||||
tenant_id: str,
|
||||
*,
|
||||
created_at: datetime.datetime | None = None,
|
||||
status: WorkflowExecutionStatus = WorkflowExecutionStatus.SUCCEEDED,
|
||||
run_type: WorkflowType = WorkflowType.WORKFLOW,
|
||||
) -> str:
|
||||
run_id = str(uuid4())
|
||||
with self.session_maker.begin() as session:
|
||||
session.add(
|
||||
WorkflowRun(
|
||||
id=run_id,
|
||||
tenant_id=tenant_id,
|
||||
app_id=str(uuid4()),
|
||||
workflow_id=str(uuid4()),
|
||||
type=run_type,
|
||||
triggered_from=WorkflowRunTriggeredFrom.APP_RUN,
|
||||
version="1",
|
||||
graph="{}",
|
||||
inputs="{}",
|
||||
status=status,
|
||||
outputs="{}",
|
||||
error=None,
|
||||
elapsed_time=0,
|
||||
total_tokens=0,
|
||||
total_steps=1,
|
||||
created_by_role=CreatorUserRole.ACCOUNT,
|
||||
created_by=str(uuid4()),
|
||||
created_at=created_at or self.end_before - datetime.timedelta(days=1),
|
||||
)
|
||||
)
|
||||
return run_id
|
||||
|
||||
def add_node(self, run_id: str, tenant_id: str, *, index: int) -> None:
|
||||
with self.session_maker.begin() as session:
|
||||
session.add(
|
||||
WorkflowNodeExecutionModel(
|
||||
id=str(uuid4()),
|
||||
tenant_id=tenant_id,
|
||||
app_id=str(uuid4()),
|
||||
workflow_id=str(uuid4()),
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
workflow_run_id=run_id,
|
||||
index=index,
|
||||
predecessor_node_id=None,
|
||||
node_execution_id=None,
|
||||
node_id=f"node-{index}",
|
||||
node_type="start",
|
||||
title="Start",
|
||||
inputs="{}",
|
||||
process_data="{}",
|
||||
outputs="{}",
|
||||
status=WorkflowNodeExecutionStatus.SUCCEEDED,
|
||||
error=None,
|
||||
elapsed_time=0,
|
||||
execution_metadata="{}",
|
||||
created_by_role=CreatorUserRole.ACCOUNT,
|
||||
created_by=str(uuid4()),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _delete_summary(
|
||||
@@ -86,100 +172,178 @@ def test_archive_tenant_id_parser_keeps_omitted_scope_unset():
|
||||
assert retention._parse_comma_separated_ids(None, param_name="tenant-ids") is None
|
||||
|
||||
|
||||
def test_resolve_archive_tenant_ids_from_plan_uses_explicit_sessions(monkeypatch):
|
||||
end_before = datetime.datetime(2025, 4, 1, tzinfo=datetime.UTC)
|
||||
sessions = [MagicMock(name="session-a"), MagicMock(name="session-b")]
|
||||
session_maker = MagicMock(side_effect=[_session_context(sessions[0]), _session_context(sessions[1])])
|
||||
calls = []
|
||||
@pytest.fixture
|
||||
def archive_db(sqlite_engine: Engine) -> ArchiveDatabase:
|
||||
"""Create only the workflow tables used by archive planning."""
|
||||
|
||||
def get_candidate_tenants(session, prefix, *, start_from, end_before):
|
||||
calls.append((session, prefix, start_from, end_before))
|
||||
return [f"{prefix}-paid", f"{prefix}-free"]
|
||||
TypeBase.metadata.create_all(
|
||||
sqlite_engine,
|
||||
tables=[WorkflowRun.__table__, WorkflowNodeExecutionModel.__table__],
|
||||
)
|
||||
return ArchiveDatabase(
|
||||
session_maker=sessionmaker(bind=sqlite_engine, expire_on_commit=False),
|
||||
end_before=datetime.datetime(2025, 4, 1, tzinfo=datetime.UTC),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(retention, "_get_archive_candidate_tenant_ids_by_prefix", get_candidate_tenants)
|
||||
|
||||
def _tenant_id(prefix: str, suffix: int) -> str:
|
||||
return f"{prefix}{suffix:07x}-0000-0000-0000-000000000000"
|
||||
|
||||
|
||||
def test_resolve_archive_tenant_ids_from_plan_uses_fresh_real_sessions(
|
||||
archive_db: ArchiveDatabase, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
paid_a = _tenant_id("a", 1)
|
||||
free_a = _tenant_id("a", 2)
|
||||
paid_b = _tenant_id("b", 1)
|
||||
free_b = _tenant_id("b", 2)
|
||||
for tenant_id in (paid_a, free_a, paid_b, free_b):
|
||||
archive_db.add_run(tenant_id)
|
||||
|
||||
# These decoys verify the real candidate query's time, status, and type filters.
|
||||
archive_db.add_run(
|
||||
_tenant_id("a", 3),
|
||||
created_at=archive_db.end_before + datetime.timedelta(seconds=1),
|
||||
)
|
||||
archive_db.add_run(_tenant_id("a", 4), status=WorkflowExecutionStatus.RUNNING)
|
||||
archive_db.add_run(_tenant_id("a", 5), run_type=WorkflowType.CHAT)
|
||||
|
||||
opened_sessions: list[Session] = []
|
||||
|
||||
def record_session(session: Session, _transaction: object, _connection: object) -> None:
|
||||
opened_sessions.append(session)
|
||||
|
||||
event.listen(archive_db.session_maker.class_, "after_begin", record_session)
|
||||
monkeypatch.setattr(
|
||||
retention,
|
||||
"_filter_paid_workflow_archive_tenant_ids",
|
||||
lambda tenant_ids: (["a-paid", "b-paid"], ["a-free", "b-free"]),
|
||||
lambda tenant_ids: ([paid_a, paid_b], sorted(set(tenant_ids) - {paid_a, paid_b})),
|
||||
)
|
||||
try:
|
||||
tenant_plan = retention._resolve_archive_tenant_ids_from_plan(
|
||||
session_maker=archive_db.session_maker,
|
||||
tenant_ids=None,
|
||||
tenant_prefixes=["a", "b"],
|
||||
start_from=None,
|
||||
end_before=archive_db.end_before,
|
||||
)
|
||||
finally:
|
||||
event.remove(archive_db.session_maker.class_, "after_begin", record_session)
|
||||
|
||||
tenant_plan = retention._resolve_archive_tenant_ids_from_plan(
|
||||
session_maker=session_maker,
|
||||
tenant_ids=None,
|
||||
tenant_prefixes=["a", "b"],
|
||||
start_from=None,
|
||||
end_before=end_before,
|
||||
)
|
||||
|
||||
assert tenant_plan["archive_tenant_ids"] == ["a-paid", "b-paid"]
|
||||
assert tenant_plan["paid_tenant_ids"] == ["a-paid", "b-paid"]
|
||||
assert tenant_plan["unpaid_tenant_ids"] == ["a-free", "b-free"]
|
||||
assert calls == [
|
||||
(sessions[0], "a", None, end_before),
|
||||
(sessions[1], "b", None, end_before),
|
||||
]
|
||||
assert tenant_plan == {
|
||||
"archive_tenant_ids": [paid_a, paid_b],
|
||||
"paid_tenant_ids": [paid_a, paid_b],
|
||||
"unpaid_tenant_ids": [free_a, free_b],
|
||||
}
|
||||
assert len(opened_sessions) == 2
|
||||
assert opened_sessions[0] is not opened_sessions[1]
|
||||
|
||||
|
||||
def test_safe_remove_scoped_session_discards_registry_and_disposes_after_remove_error(monkeypatch):
|
||||
fake_db = MagicMock()
|
||||
fake_db.session.remove.side_effect = RuntimeError("server closed the connection unexpectedly")
|
||||
monkeypatch.setattr(retention, "db", fake_db)
|
||||
def test_safe_remove_scoped_session_recovers_from_real_closed_connection(
|
||||
sqlite_engine: Engine,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
maker = sessionmaker(bind=sqlite_engine)
|
||||
registry = scoped_session(maker)
|
||||
registry().execute(text("select 1"))
|
||||
dbapi_connection = registry().connection().connection.dbapi_connection
|
||||
assert dbapi_connection is not None
|
||||
dbapi_connection.close()
|
||||
monkeypatch.setattr(retention, "db", SimpleNamespace(session=registry, engine=sqlite_engine))
|
||||
|
||||
retention._safe_remove_scoped_session("archive workflow run command")
|
||||
with caplog.at_level(logging.WARNING, logger="commands.retention"):
|
||||
retention._safe_remove_scoped_session("archive workflow run command")
|
||||
|
||||
fake_db.session.remove.assert_called_once()
|
||||
fake_db.session.registry.clear.assert_called_once()
|
||||
fake_db.engine.dispose.assert_called_once()
|
||||
assert not registry.registry.has()
|
||||
assert any("Ignoring DB scoped-session cleanup error" in message for message in caplog.messages)
|
||||
|
||||
|
||||
def test_archive_command_db_retry_retries_retryable_db_disconnect(monkeypatch):
|
||||
operation = MagicMock(side_effect=[_db_disconnect_error(), "ok"])
|
||||
sleep = MagicMock()
|
||||
def test_archive_command_db_retry_retries_retryable_db_disconnect(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
attempts = iter([_db_disconnect_error(), "ok"])
|
||||
sleep = Mock()
|
||||
monkeypatch.setattr("services.retention.workflow_run.db_retry.time.sleep", sleep)
|
||||
|
||||
result = retention._run_archive_command_db_retry("archive plan", operation)
|
||||
def operation() -> str:
|
||||
result = next(attempts)
|
||||
if isinstance(result, Exception):
|
||||
raise result
|
||||
return result
|
||||
|
||||
assert result == "ok"
|
||||
assert operation.call_count == 2
|
||||
assert retention._run_archive_command_db_retry("archive plan", operation) == "ok"
|
||||
sleep.assert_called_once_with(1.0)
|
||||
|
||||
|
||||
def test_archive_plan_prefix_stats_retries_count_query_with_fresh_session(monkeypatch):
|
||||
end_before = datetime.datetime(2025, 4, 1, tzinfo=datetime.UTC)
|
||||
sessions = [MagicMock(name="session-1"), MagicMock(name="session-2")]
|
||||
sessions[0].scalar.side_effect = _db_disconnect_error()
|
||||
sessions[1].scalar.side_effect = [7, 9]
|
||||
session_maker = MagicMock(side_effect=[_session_context(sessions[0]), _session_context(sessions[1])])
|
||||
sleep = MagicMock()
|
||||
def test_archive_plan_prefix_stats_retries_with_fresh_session_and_real_counts(
|
||||
archive_db: ArchiveDatabase, sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
tenant_id = _tenant_id("a", 1)
|
||||
run_ids = [archive_db.add_run(tenant_id) for _ in range(7)]
|
||||
for index in range(9):
|
||||
archive_db.add_node(run_ids[index % len(run_ids)], tenant_id, index=index)
|
||||
|
||||
monkeypatch.setattr(
|
||||
retention,
|
||||
"_get_archive_candidate_tenant_ids_by_prefix",
|
||||
lambda session, prefix, *, start_from, end_before: [f"{prefix}-tenant"],
|
||||
# Decoys outside the selected prefix and archive window must not affect counts.
|
||||
decoy_run_id = archive_db.add_run(_tenant_id("b", 1))
|
||||
archive_db.add_node(decoy_run_id, _tenant_id("b", 1), index=99)
|
||||
archive_db.add_run(
|
||||
tenant_id,
|
||||
created_at=archive_db.end_before + datetime.timedelta(seconds=1),
|
||||
)
|
||||
|
||||
fail_next_query = True
|
||||
|
||||
def disconnect_once(
|
||||
_connection: object,
|
||||
_cursor: object,
|
||||
_statement: str,
|
||||
_parameters: object,
|
||||
_context: object,
|
||||
_executemany: bool,
|
||||
) -> None:
|
||||
nonlocal fail_next_query
|
||||
if fail_next_query:
|
||||
fail_next_query = False
|
||||
raise _db_disconnect_error()
|
||||
|
||||
opened_sessions: list[Session] = []
|
||||
|
||||
def record_session(session: Session, _transaction: object, _connection: object) -> None:
|
||||
opened_sessions.append(session)
|
||||
|
||||
sleep = Mock()
|
||||
monkeypatch.setattr("services.retention.workflow_run.db_retry.time.sleep", sleep)
|
||||
event.listen(sqlite_engine, "before_cursor_execute", disconnect_once)
|
||||
event.listen(archive_db.session_maker.class_, "after_begin", record_session)
|
||||
try:
|
||||
stats = retention._get_archive_plan_prefix_stats(
|
||||
archive_db.session_maker,
|
||||
"a",
|
||||
start_from=None,
|
||||
end_before=archive_db.end_before,
|
||||
)
|
||||
finally:
|
||||
event.remove(sqlite_engine, "before_cursor_execute", disconnect_once)
|
||||
event.remove(archive_db.session_maker.class_, "after_begin", record_session)
|
||||
|
||||
stats = retention._get_archive_plan_prefix_stats(
|
||||
session_maker,
|
||||
"a",
|
||||
start_from=None,
|
||||
end_before=end_before,
|
||||
)
|
||||
|
||||
assert stats["tenant_ids"] == ["a-tenant"]
|
||||
assert stats["workflow_runs"] == 7
|
||||
assert stats["workflow_node_executions"] == 9
|
||||
assert session_maker.call_count == 2
|
||||
assert stats == {
|
||||
"tenant_ids": [tenant_id],
|
||||
"workflow_runs": 7,
|
||||
"workflow_node_executions": 9,
|
||||
}
|
||||
assert len(opened_sessions) == 2
|
||||
assert opened_sessions[0] is not opened_sessions[1]
|
||||
sleep.assert_called_once_with(1.0)
|
||||
|
||||
|
||||
def test_archive_workflow_runs_raises_click_exception_when_tenant_plan_fails(monkeypatch):
|
||||
fake_db = MagicMock()
|
||||
monkeypatch.setattr(retention, "db", fake_db)
|
||||
def test_archive_workflow_runs_raises_click_exception_when_tenant_plan_fails(
|
||||
sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
registry = scoped_session(sessionmaker(bind=sqlite_engine))
|
||||
monkeypatch.setattr(retention, "db", SimpleNamespace(engine=sqlite_engine, session=registry))
|
||||
monkeypatch.setattr(
|
||||
retention,
|
||||
"_resolve_archive_tenant_ids_from_plan",
|
||||
MagicMock(side_effect=RuntimeError("tenant plan failed")),
|
||||
Mock(side_effect=RuntimeError("tenant plan failed")),
|
||||
)
|
||||
|
||||
with pytest.raises(click.ClickException, match="Failed to resolve workflow archive tenant plan"):
|
||||
|
||||
Reference in New Issue
Block a user