mirror of
https://github.com/langgenius/dify.git
synced 2026-11-12 06:11:12 +08:00
640 lines
31 KiB
Python
640 lines
31 KiB
Python
"""Tests for Celery SSL configuration."""
|
|
|
|
import ssl
|
|
from collections.abc import Iterator
|
|
from datetime import UTC, datetime, timedelta
|
|
from importlib import import_module
|
|
from queue import Empty
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, Mock, patch
|
|
from zoneinfo import ZoneInfo
|
|
|
|
import pytest
|
|
from celery import Celery
|
|
from celery.contrib.testing.app import TestApp, setup_default_app
|
|
from celery.exceptions import Retry
|
|
|
|
from configs.feature import LoggingConfig
|
|
from dify_app import DifyApp
|
|
from enums import DeploymentEdition
|
|
from extensions.ext_celery import (
|
|
_enqueue_initial_community_telemetry_heartbeat,
|
|
get_celery_broker_transport_options,
|
|
init_app,
|
|
)
|
|
from schedule.collect_agent_sandbox_usage import collect_agent_sandbox_usage
|
|
from services.knowledge_fs.background_contract import OPERATION_QUEUES, OPERATION_TASK
|
|
from services.knowledge_fs.lifecycle_readiness import KnowledgeFSLifecycleWorkerReadiness
|
|
from tasks.application_trace_task import export_application_trace
|
|
from tasks.trace_delivery_task import deliver_trace_event, retry_trace_export
|
|
from tasks.workflow_draft_coordination_task import (
|
|
process_workflow_draft_coordination,
|
|
sweep_workflow_draft_coordination,
|
|
)
|
|
from tests.helpers.unit.config_override import apply_config_overrides, config_overrides_context
|
|
from tests.helpers.unit.worker_entrypoint import worker_configuration
|
|
from tests.unit_tests.celery_fixtures import preserve_celery_apps
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("configured", "iana_name"),
|
|
[
|
|
("utc", "UTC"),
|
|
("asia/shanghai", "Asia/Shanghai"),
|
|
],
|
|
)
|
|
def test_init_app_uses_validated_logging_timezone(configured: str, iana_name: str) -> None:
|
|
logging_config = LoggingConfig(LOG_TZ=configured)
|
|
assert logging_config.LOG_TZ is not None
|
|
|
|
with (
|
|
setup_default_app(TestApp()),
|
|
config_overrides_context(
|
|
CELERY_BROKER_URL="memory://",
|
|
CELERY_BACKEND="disabled",
|
|
DISABLE_TELEMETRY=True,
|
|
LOG_TZ=logging_config.LOG_TZ,
|
|
),
|
|
):
|
|
celery_app = init_app(DifyApp(__name__))
|
|
try:
|
|
expected_timezone = ZoneInfo(iana_name)
|
|
for instant in (
|
|
datetime(2026, 1, 15, 12, tzinfo=UTC),
|
|
datetime(2026, 7, 15, 12, tzinfo=UTC),
|
|
):
|
|
assert (
|
|
instant.astimezone(celery_app.conf.timezone).utcoffset()
|
|
== instant.astimezone(expected_timezone).utcoffset()
|
|
)
|
|
finally:
|
|
celery_app.close()
|
|
|
|
|
|
@pytest.fixture
|
|
def tracing_celery(request: pytest.FixtureRequest) -> Iterator[Celery]:
|
|
with (
|
|
setup_default_app(TestApp()),
|
|
config_overrides_context(
|
|
CELERY_BROKER_URL="memory://",
|
|
CELERY_BACKEND="disabled",
|
|
OPS_TRACE_QUEUE=request.param,
|
|
OPS_AGENT_TRACE_ENABLED=True,
|
|
ENABLE_APPLICATION_TRACE_RECONCILE_TASK=True,
|
|
ENABLE_TRACE_PAYLOAD_SWEEP_TASK=True,
|
|
DISABLE_TELEMETRY=True,
|
|
),
|
|
):
|
|
app = init_app(DifyApp(__name__))
|
|
try:
|
|
yield app
|
|
finally:
|
|
app.close()
|
|
|
|
|
|
@pytest.mark.parametrize("tracing_celery", ["ops_trace_v2", "ops_trace", "tracing-upgrade"], indirect=True)
|
|
def test_tracing_dispatch_and_beat_use_configured_queue(tracing_celery: Celery) -> None:
|
|
queue = tracing_celery.conf.task_routes["tasks.trace_delivery_task.*"]["queue"]
|
|
calls = {
|
|
deliver_trace_event.name: ({"event": "{}"},),
|
|
retry_trace_export.name: ("trace_payloads/job.json",),
|
|
export_application_trace.name: ("run-1",),
|
|
tracing_celery.conf.beat_schedule["sweep_trace_payloads"]["task"]: (),
|
|
tracing_celery.conf.beat_schedule["reconcile_application_traces"]["task"]: (),
|
|
collect_agent_sandbox_usage.name: (),
|
|
}
|
|
with tracing_celery.connection_for_write() as connection:
|
|
if queue != "ops_trace":
|
|
tracing_celery.send_task("tasks.ops_trace_task.process_trace_tasks", args=({},), queue="ops_trace")
|
|
for name, args in calls.items():
|
|
tracing_celery.tasks[name].apply_async(args=args)
|
|
with connection.SimpleQueue(queue) as inbox:
|
|
for name, args in calls.items():
|
|
message = inbox.get(block=False)
|
|
assert message.headers["task"] == name
|
|
assert message.payload[0] == list(args)
|
|
message.ack()
|
|
with pytest.raises(Empty):
|
|
inbox.get(block=False)
|
|
if queue != "ops_trace":
|
|
with connection.SimpleQueue("ops_trace") as legacy:
|
|
message = legacy.get(block=False)
|
|
assert message.headers["task"] == "tasks.ops_trace_task.process_trace_tasks"
|
|
message.ack()
|
|
|
|
|
|
@pytest.mark.parametrize("tracing_celery", ["ops_trace"], indirect=True)
|
|
def test_redelivery_stays_on_original_queue_during_switchback(tracing_celery: Celery) -> None:
|
|
task = tracing_celery.tasks["tasks.trace_delivery_task.deliver_trace_event"]
|
|
task.push_request(
|
|
id="trace-redelivery",
|
|
args=({"event": "{}"},),
|
|
kwargs={},
|
|
retries=0,
|
|
called_directly=False,
|
|
delivery_info={"exchange": "", "routing_key": "ops_trace_v2"},
|
|
)
|
|
try:
|
|
with pytest.raises(Retry):
|
|
task.retry(countdown=5)
|
|
with tracing_celery.connection_for_read() as connection:
|
|
with connection.SimpleQueue("ops_trace_v2") as original:
|
|
message = original.get(block=False)
|
|
assert message.headers["task"] == task.name
|
|
assert message.headers["retries"] == 1
|
|
message.ack()
|
|
with connection.SimpleQueue("ops_trace") as restored:
|
|
with pytest.raises(Empty):
|
|
restored.get(block=False)
|
|
finally:
|
|
task.pop_request()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"edition", [DeploymentEdition.COMMUNITY, DeploymentEdition.ENTERPRISE, DeploymentEdition.CLOUD]
|
|
)
|
|
def test_draft_recovery_dispatch_reaches_default_workers_without_builder(edition: DeploymentEdition) -> None:
|
|
selected = worker_configuration({"DEPLOYMENT_EDITION": edition.value})
|
|
assert selected.returncode == 0, selected.stderr
|
|
subscribed_queues = selected.stdout.splitlines()[-1].split(",")
|
|
with (
|
|
setup_default_app(TestApp()),
|
|
config_overrides_context(
|
|
CELERY_BROKER_URL="memory://",
|
|
CELERY_BACKEND="disabled",
|
|
DEPLOYMENT_EDITION=edition,
|
|
DIFY_BUILDER_ENABLED=False,
|
|
ENABLE_DIFY_BUILDER_RECONCILER=False,
|
|
ENABLE_WORKFLOW_DRAFT_COORDINATION_TASK=True,
|
|
DISABLE_TELEMETRY=True,
|
|
),
|
|
):
|
|
app = init_app(DifyApp(__name__))
|
|
try:
|
|
process_name = process_workflow_draft_coordination.name
|
|
sweep_name = sweep_workflow_draft_coordination.name
|
|
schedule = app.conf.beat_schedule["workflow_draft_coordination_sweeper"]
|
|
assert schedule["task"] == sweep_name
|
|
app.tasks[process_name].apply_async(args=("accepted-draft",))
|
|
app.tasks[sweep_name].apply_async()
|
|
app.tasks[sweep_name].apply_async(**schedule["options"])
|
|
deliveries = {
|
|
"workflow_storage": [(process_name, ["accepted-draft"])],
|
|
"schedule_executor": [(sweep_name, []), (sweep_name, [])],
|
|
}
|
|
with app.connection_for_read() as connection:
|
|
for queue, expected in deliveries.items():
|
|
assert queue in subscribed_queues
|
|
with connection.SimpleQueue(queue) as inbox:
|
|
for task_name, args in expected:
|
|
message = inbox.get(block=False)
|
|
assert message.headers["task"] == task_name
|
|
assert message.payload[0] == args
|
|
message.ack()
|
|
with pytest.raises(Empty):
|
|
inbox.get(block=False)
|
|
finally:
|
|
app.close()
|
|
|
|
|
|
class TestCelerySslCases:
|
|
@staticmethod
|
|
@pytest.fixture(autouse=True)
|
|
def restore_celery_app_state() -> Iterator[None]:
|
|
# init_app replaces Celery's process-wide current/default app. Restore both
|
|
# so later shared tasks do not enter this module's Flask app context.
|
|
with setup_default_app(TestApp()):
|
|
yield
|
|
|
|
class TestCelerySSLConfiguration:
|
|
"""Test suite for Celery SSL configuration."""
|
|
|
|
def test_get_celery_broker_transport_options_includes_global_keyprefix_for_redis(self):
|
|
mock_config = MagicMock()
|
|
mock_config.CELERY_USE_SENTINEL = False
|
|
mock_config.REDIS_KEY_PREFIX = "enterprise-a"
|
|
with patch("extensions.ext_celery.dify_config", mock_config):
|
|
from extensions.ext_celery import get_celery_broker_transport_options
|
|
|
|
result = get_celery_broker_transport_options()
|
|
assert result["global_keyprefix"] == "enterprise-a:"
|
|
|
|
def test_get_celery_broker_transport_options_omits_global_keyprefix_when_prefix_empty(self):
|
|
mock_config = MagicMock()
|
|
mock_config.CELERY_USE_SENTINEL = False
|
|
mock_config.REDIS_KEY_PREFIX = " "
|
|
with patch("extensions.ext_celery.dify_config", mock_config):
|
|
from extensions.ext_celery import get_celery_broker_transport_options
|
|
|
|
result = get_celery_broker_transport_options()
|
|
assert "global_keyprefix" not in result
|
|
|
|
def test_get_celery_broker_transport_options_keeps_sentinel_and_adds_global_keyprefix(self):
|
|
mock_config = MagicMock()
|
|
mock_config.CELERY_USE_SENTINEL = True
|
|
mock_config.CELERY_SENTINEL_MASTER_NAME = "mymaster"
|
|
mock_config.CELERY_SENTINEL_SOCKET_TIMEOUT = 0.1
|
|
mock_config.CELERY_SENTINEL_PASSWORD = "secret"
|
|
mock_config.REDIS_KEY_PREFIX = "enterprise-a"
|
|
with patch("extensions.ext_celery.dify_config", mock_config):
|
|
from extensions.ext_celery import get_celery_broker_transport_options
|
|
|
|
result = get_celery_broker_transport_options()
|
|
assert result["master_name"] == "mymaster"
|
|
assert result["sentinel_kwargs"]["password"] == "secret"
|
|
assert result["global_keyprefix"] == "enterprise-a:"
|
|
|
|
def test_get_celery_ssl_options_when_ssl_disabled(self):
|
|
"""Test SSL options when BROKER_USE_SSL is False."""
|
|
from configs import DifyConfig
|
|
|
|
dify_config = DifyConfig(CELERY_BROKER_URL="redis://localhost:6379/0")
|
|
with patch("extensions.ext_celery.dify_config", dify_config):
|
|
from extensions.ext_celery import get_celery_ssl_options
|
|
|
|
result = get_celery_ssl_options()
|
|
assert result is None
|
|
|
|
def test_get_celery_ssl_options_when_broker_not_redis(self):
|
|
"""Test SSL options when broker is not Redis."""
|
|
mock_config = MagicMock()
|
|
mock_config.CELERY_BROKER_URL = "amqp://localhost:5672"
|
|
with patch("extensions.ext_celery.dify_config", mock_config):
|
|
from extensions.ext_celery import get_celery_ssl_options
|
|
|
|
result = get_celery_ssl_options()
|
|
assert result is None
|
|
|
|
def test_get_celery_ssl_options_with_cert_none(self):
|
|
"""Test SSL options with CERT_NONE requirement."""
|
|
mock_config = MagicMock()
|
|
mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0"
|
|
mock_config.REDIS_SSL_CERT_REQS = "CERT_NONE"
|
|
mock_config.REDIS_SSL_CA_CERTS = None
|
|
mock_config.REDIS_SSL_CERTFILE = None
|
|
mock_config.REDIS_SSL_KEYFILE = None
|
|
with patch("extensions.ext_celery.dify_config", mock_config):
|
|
from extensions.ext_celery import get_celery_ssl_options
|
|
|
|
result = get_celery_ssl_options()
|
|
assert result is not None
|
|
assert result["ssl_cert_reqs"] == ssl.CERT_NONE
|
|
assert result["ssl_ca_certs"] is None
|
|
assert result["ssl_certfile"] is None
|
|
assert result["ssl_keyfile"] is None
|
|
|
|
def test_get_celery_ssl_options_with_cert_required(self):
|
|
"""Test SSL options with CERT_REQUIRED and certificates."""
|
|
mock_config = MagicMock()
|
|
mock_config.CELERY_BROKER_URL = "rediss://localhost:6380/0"
|
|
mock_config.REDIS_SSL_CERT_REQS = "CERT_REQUIRED"
|
|
mock_config.REDIS_SSL_CA_CERTS = "/path/to/ca.crt"
|
|
mock_config.REDIS_SSL_CERTFILE = "/path/to/client.crt"
|
|
mock_config.REDIS_SSL_KEYFILE = "/path/to/client.key"
|
|
with patch("extensions.ext_celery.dify_config", mock_config):
|
|
from extensions.ext_celery import get_celery_ssl_options
|
|
|
|
result = get_celery_ssl_options()
|
|
assert result is not None
|
|
assert result["ssl_cert_reqs"] == ssl.CERT_REQUIRED
|
|
assert result["ssl_ca_certs"] == "/path/to/ca.crt"
|
|
assert result["ssl_certfile"] == "/path/to/client.crt"
|
|
assert result["ssl_keyfile"] == "/path/to/client.key"
|
|
|
|
def test_get_celery_ssl_options_with_cert_optional(self):
|
|
"""Test SSL options with CERT_OPTIONAL requirement."""
|
|
mock_config = MagicMock()
|
|
mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0"
|
|
mock_config.REDIS_SSL_CERT_REQS = "CERT_OPTIONAL"
|
|
mock_config.REDIS_SSL_CA_CERTS = "/path/to/ca.crt"
|
|
mock_config.REDIS_SSL_CERTFILE = None
|
|
mock_config.REDIS_SSL_KEYFILE = None
|
|
with patch("extensions.ext_celery.dify_config", mock_config):
|
|
from extensions.ext_celery import get_celery_ssl_options
|
|
|
|
result = get_celery_ssl_options()
|
|
assert result is not None
|
|
assert result["ssl_cert_reqs"] == ssl.CERT_OPTIONAL
|
|
assert result["ssl_ca_certs"] == "/path/to/ca.crt"
|
|
|
|
def test_get_celery_ssl_options_with_invalid_cert_reqs(self):
|
|
"""Test SSL options with invalid cert requirement defaults to CERT_NONE."""
|
|
mock_config = MagicMock()
|
|
mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0"
|
|
mock_config.REDIS_SSL_CERT_REQS = "INVALID_VALUE"
|
|
mock_config.REDIS_SSL_CA_CERTS = None
|
|
mock_config.REDIS_SSL_CERTFILE = None
|
|
mock_config.REDIS_SSL_KEYFILE = None
|
|
with patch("extensions.ext_celery.dify_config", mock_config):
|
|
from extensions.ext_celery import get_celery_ssl_options
|
|
|
|
result = get_celery_ssl_options()
|
|
assert result is not None
|
|
assert result["ssl_cert_reqs"] == ssl.CERT_NONE
|
|
|
|
def test_celery_init_applies_ssl_to_broker_and_backend(self):
|
|
"""Test that SSL options are applied to both broker and backend when using Redis."""
|
|
mock_config = MagicMock()
|
|
mock_config.HUMAN_INPUT_TIMEOUT_TASK_INTERVAL = 1
|
|
mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0"
|
|
mock_config.CELERY_BACKEND = "redis"
|
|
mock_config.CELERY_RESULT_BACKEND = "redis://localhost:6379/0"
|
|
mock_config.REDIS_SSL_CERT_REQS = "CERT_NONE"
|
|
mock_config.REDIS_SSL_CA_CERTS = None
|
|
mock_config.REDIS_SSL_CERTFILE = None
|
|
mock_config.REDIS_SSL_KEYFILE = None
|
|
mock_config.CELERY_USE_SENTINEL = False
|
|
mock_config.LOG_FORMAT = "%(message)s"
|
|
mock_config.LOG_TZ = "UTC"
|
|
mock_config.LOG_FILE = None
|
|
mock_config.CELERY_BEAT_SCHEDULER_TIME = 1
|
|
mock_config.AGENT_SANDBOX_METERING_ENABLED = False
|
|
mock_config.ENABLE_CONVERSATION_CLEANUP_TASK = False
|
|
mock_config.ENABLE_DIFY_BUILDER_RECONCILER = False
|
|
mock_config.ENABLE_WORKFLOW_DRAFT_COORDINATION_TASK = True
|
|
mock_config.WORKFLOW_DRAFT_COORDINATION_INTERVAL = 60
|
|
mock_config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5
|
|
mock_config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK = False
|
|
mock_config.ENABLE_CLEAN_UNUSED_DATASETS_TASK = False
|
|
mock_config.ENABLE_CREATE_TIDB_SERVERLESS_TASK = False
|
|
mock_config.ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK = False
|
|
mock_config.ENABLE_CLEAN_MESSAGES = False
|
|
mock_config.ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK = False
|
|
mock_config.ENABLE_DATASETS_QUEUE_MONITOR = False
|
|
mock_config.ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK = False
|
|
mock_config.ENABLE_WORKFLOW_SCHEDULE_POLLER_TASK = False
|
|
mock_config.WORKFLOW_SCHEDULE_POLLER_INTERVAL = 1
|
|
mock_config.WORKFLOW_SCHEDULE_POLLER_BATCH_SIZE = 100
|
|
mock_config.WORKFLOW_SCHEDULE_MAX_DISPATCH_PER_TICK = 0
|
|
mock_config.ENABLE_TRIGGER_PROVIDER_REFRESH_TASK = False
|
|
mock_config.TRIGGER_PROVIDER_REFRESH_INTERVAL = 15
|
|
mock_config.ENABLE_API_TOKEN_LAST_USED_UPDATE_TASK = False
|
|
mock_config.API_TOKEN_LAST_USED_UPDATE_INTERVAL = 30
|
|
with (
|
|
patch("extensions.ext_celery.dify_config", mock_config),
|
|
patch(
|
|
"services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness",
|
|
return_value=KnowledgeFSLifecycleWorkerReadiness(ready=False, blockers=()),
|
|
),
|
|
):
|
|
from dify_app import DifyApp
|
|
from extensions.ext_celery import init_app
|
|
|
|
app = DifyApp(__name__)
|
|
celery_app = init_app(app)
|
|
assert "broker_use_ssl" in celery_app.conf
|
|
assert celery_app.conf["broker_use_ssl"] is not None
|
|
assert celery_app.conf["broker_use_ssl"]["ssl_cert_reqs"] == ssl.CERT_NONE
|
|
assert "redis_backend_use_ssl" in celery_app.conf
|
|
assert celery_app.conf["redis_backend_use_ssl"] is not None
|
|
|
|
def test_celery_init_registers_required_agent_and_conversation_tasks(self):
|
|
mock_config = MagicMock()
|
|
mock_config.BROKER_USE_SSL = False
|
|
mock_config.REDIS_KEY_PREFIX = "enterprise-a"
|
|
mock_config.HUMAN_INPUT_TIMEOUT_TASK_INTERVAL = 1
|
|
mock_config.CELERY_BROKER_URL = "redis://localhost:6379/0"
|
|
mock_config.CELERY_BACKEND = "redis"
|
|
mock_config.CELERY_RESULT_BACKEND = "redis://localhost:6379/0"
|
|
mock_config.CELERY_USE_SENTINEL = False
|
|
mock_config.LOG_FORMAT = "%(message)s"
|
|
mock_config.LOG_TZ = "UTC"
|
|
mock_config.LOG_FILE = None
|
|
mock_config.CELERY_TASK_ANNOTATIONS = {}
|
|
mock_config.CELERY_BEAT_SCHEDULER_TIME = 1
|
|
mock_config.AGENT_SANDBOX_METERING_ENABLED = False
|
|
mock_config.ENABLE_CONVERSATION_CLEANUP_TASK = True
|
|
mock_config.ENABLE_DIFY_BUILDER_RECONCILER = False
|
|
mock_config.ENABLE_WORKFLOW_DRAFT_COORDINATION_TASK = True
|
|
mock_config.WORKFLOW_DRAFT_COORDINATION_INTERVAL = 60
|
|
mock_config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5
|
|
mock_config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK = False
|
|
mock_config.ENABLE_CLEAN_UNUSED_DATASETS_TASK = False
|
|
mock_config.ENABLE_CREATE_TIDB_SERVERLESS_TASK = False
|
|
mock_config.ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK = False
|
|
mock_config.ENABLE_CLEAN_MESSAGES = False
|
|
mock_config.ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK = False
|
|
mock_config.ENABLE_DATASETS_QUEUE_MONITOR = False
|
|
mock_config.ENABLE_HUMAN_INPUT_TIMEOUT_TASK = False
|
|
mock_config.ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK = False
|
|
mock_config.MARKETPLACE_ENABLED = False
|
|
mock_config.WORKFLOW_LOG_CLEANUP_ENABLED = False
|
|
mock_config.ENABLE_WORKFLOW_RUN_CLEANUP_TASK = False
|
|
mock_config.ENABLE_WORKFLOW_SCHEDULE_POLLER_TASK = False
|
|
mock_config.WORKFLOW_SCHEDULE_POLLER_INTERVAL = 1
|
|
mock_config.ENABLE_TRIGGER_PROVIDER_REFRESH_TASK = False
|
|
mock_config.TRIGGER_PROVIDER_REFRESH_INTERVAL = 15
|
|
mock_config.ENABLE_API_TOKEN_LAST_USED_UPDATE_TASK = False
|
|
mock_config.API_TOKEN_LAST_USED_UPDATE_INTERVAL = 30
|
|
mock_config.DEPLOYMENT_EDITION = DeploymentEdition.COMMUNITY
|
|
mock_config.ENTERPRISE_TELEMETRY_ENABLED = False
|
|
with (
|
|
patch("extensions.ext_celery.dify_config", mock_config),
|
|
patch(
|
|
"services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness",
|
|
return_value=KnowledgeFSLifecycleWorkerReadiness(ready=False, blockers=()),
|
|
),
|
|
):
|
|
from dify_app import DifyApp
|
|
from extensions.ext_celery import init_app
|
|
|
|
app = DifyApp(__name__)
|
|
celery_app = init_app(app)
|
|
assert celery_app.conf["broker_transport_options"]["global_keyprefix"] == "enterprise-a:"
|
|
assert celery_app.conf["result_backend_transport_options"]["global_keyprefix"] == "enterprise-a:"
|
|
assert "tasks.collect_agent_resources_task" in celery_app.conf["imports"]
|
|
assert "tasks.delete_conversation_task" in celery_app.conf["imports"]
|
|
assert "tasks.dify_builder_advance_task" in celery_app.conf["imports"]
|
|
assert "tasks.workflow_run_archive_deletion_tasks" in celery_app.conf["imports"]
|
|
import_module("tasks.workflow_run_archive_deletion_tasks")
|
|
task_name = "tasks.workflow_run_archive_deletion_tasks.delete_workflow_run_archive_task"
|
|
archive_deletion = celery_app.tasks[task_name]
|
|
assert archive_deletion.queue == "workflow_archive"
|
|
assert archive_deletion.acks_late is True
|
|
assert archive_deletion.reject_on_worker_lost is True
|
|
assert archive_deletion.autoretry_for == (Exception,)
|
|
assert archive_deletion.retry_kwargs["max_retries"] == 5
|
|
assert (
|
|
celery_app.conf["beat_schedule"]["conversation_cleanup_sweeper"]["task"]
|
|
== "tasks.delete_conversation_task.sweep_deleted_conversations"
|
|
)
|
|
draft_sweeper = celery_app.conf["beat_schedule"]["workflow_draft_coordination_sweeper"]
|
|
assert draft_sweeper["task"] == "tasks.workflow_draft_coordination_task.sweep_workflow_draft_coordination"
|
|
assert draft_sweeper["schedule"].total_seconds() == 60
|
|
assert draft_sweeper["options"] == {"queue": "schedule_executor"}
|
|
|
|
|
|
class TestCommunityTelemetryCeleryCases:
|
|
@staticmethod
|
|
def test_beat_start_enqueues_community_telemetry_heartbeat() -> None:
|
|
task = Mock()
|
|
sender = SimpleNamespace(
|
|
app=SimpleNamespace(
|
|
conf=SimpleNamespace(beat_schedule={"community_telemetry_heartbeat": {}}),
|
|
tasks={"community_telemetry.send_heartbeat": task},
|
|
)
|
|
)
|
|
|
|
_enqueue_initial_community_telemetry_heartbeat(sender)
|
|
|
|
task.apply_async.assert_called_once_with()
|
|
|
|
@staticmethod
|
|
def test_beat_start_skips_community_telemetry_when_not_scheduled() -> None:
|
|
task = Mock()
|
|
sender = SimpleNamespace(
|
|
app=SimpleNamespace(
|
|
conf=SimpleNamespace(beat_schedule={}),
|
|
tasks={"community_telemetry.send_heartbeat": task},
|
|
)
|
|
)
|
|
|
|
_enqueue_initial_community_telemetry_heartbeat(sender)
|
|
|
|
task.apply_async.assert_not_called()
|
|
|
|
@staticmethod
|
|
def test_beat_start_skips_community_telemetry_when_task_is_unavailable() -> None:
|
|
sender = SimpleNamespace(
|
|
app=SimpleNamespace(
|
|
conf=SimpleNamespace(beat_schedule={"community_telemetry_heartbeat": {}}),
|
|
tasks={},
|
|
)
|
|
)
|
|
|
|
_enqueue_initial_community_telemetry_heartbeat(sender)
|
|
|
|
|
|
class TestKnowledgeFsCeleryCases:
|
|
@staticmethod
|
|
@pytest.fixture(autouse=True)
|
|
def restore_celery_app_state() -> Iterator[None]:
|
|
# init_app replaces Celery's process-wide current/default app; restore both so
|
|
# later shared tasks still resolve against the app that registered them.
|
|
with setup_default_app(TestApp()):
|
|
yield
|
|
|
|
@staticmethod
|
|
@pytest.mark.parametrize("background_enabled", [False, True])
|
|
@pytest.mark.parametrize("vector_store", ["qdrant", "tidb_on_qdrant", "elasticsearch", "weaviate"])
|
|
def test_celery_registers_initial_source_task_when_knowledge_fs_lifecycle_is_ready(
|
|
background_enabled: bool, vector_store: str, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
# A stale deployment variable must not suppress vector cleanup for the selected provider.
|
|
monkeypatch.setenv("KNOWLEDGE_VECTOR_STORAGE", "postgres")
|
|
apply_config_overrides(
|
|
monkeypatch,
|
|
REDIS_KEY_PREFIX="test",
|
|
HUMAN_INPUT_TIMEOUT_TASK_INTERVAL=1,
|
|
CELERY_BROKER_URL="redis://localhost:6379/0",
|
|
CELERY_BACKEND="redis",
|
|
CELERY_USE_SENTINEL=False,
|
|
CELERY_BROKER_VISIBILITY_TIMEOUT=10800,
|
|
KNOWLEDGE_FS_BACKGROUND_WORKER_ENABLED=background_enabled,
|
|
VECTOR_STORE=vector_store,
|
|
KNOWLEDGE_FS_BACKGROUND_POLL_INTERVAL_SECONDS=5,
|
|
KNOWLEDGE_FS_BACKGROUND_TASK_TIMEOUT_SECONDS=7200,
|
|
LOG_FORMAT="%(message)s",
|
|
LOG_TZ="UTC",
|
|
LOG_FILE=None,
|
|
CELERY_TASK_ANNOTATIONS=dict[str, object](),
|
|
CELERY_BEAT_SCHEDULER_TIME=1,
|
|
KNOWLEDGE_FS_LIFECYCLE_POLL_INTERVAL_SECONDS=2,
|
|
ENABLE_CONVERSATION_CLEANUP_TASK=True,
|
|
CONVERSATION_CLEANUP_TASK_INTERVAL=5,
|
|
ENABLE_CLEAN_EMBEDDING_CACHE_TASK=False,
|
|
ENABLE_CLEAN_UNUSED_DATASETS_TASK=False,
|
|
ENABLE_CREATE_TIDB_SERVERLESS_TASK=False,
|
|
ENABLE_UPDATE_TIDB_SERVERLESS_STATUS_TASK=False,
|
|
ENABLE_CLEAN_MESSAGES=False,
|
|
ENABLE_MAIL_CLEAN_DOCUMENT_NOTIFY_TASK=False,
|
|
ENABLE_DATASETS_QUEUE_MONITOR=False,
|
|
ENABLE_HUMAN_INPUT_TIMEOUT_TASK=False,
|
|
ENABLE_CHECK_UPGRADABLE_PLUGIN_TASK=False,
|
|
MARKETPLACE_ENABLED=False,
|
|
WORKFLOW_LOG_CLEANUP_ENABLED=False,
|
|
ENABLE_WORKFLOW_RUN_CLEANUP_TASK=False,
|
|
ENABLE_WORKFLOW_SCHEDULE_POLLER_TASK=False,
|
|
WORKFLOW_SCHEDULE_POLLER_INTERVAL=1,
|
|
ENABLE_TRIGGER_PROVIDER_REFRESH_TASK=False,
|
|
TRIGGER_PROVIDER_REFRESH_INTERVAL=15,
|
|
ENABLE_API_TOKEN_LAST_USED_UPDATE_TASK=False,
|
|
API_TOKEN_LAST_USED_UPDATE_INTERVAL=30,
|
|
ENTERPRISE_TELEMETRY_ENABLED=False,
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness",
|
|
return_value=KnowledgeFSLifecycleWorkerReadiness(ready=True, blockers=()),
|
|
),
|
|
):
|
|
celery_app = init_app(DifyApp(__name__))
|
|
|
|
assert "tasks.knowledge_fs_initial_source_tasks" in celery_app.conf["imports"]
|
|
assert "tasks.knowledge_fs_initial_source_preview_tasks" in celery_app.conf["imports"]
|
|
assert "tasks.knowledge_fs_failed_retrieval_tasks" in celery_app.conf["imports"]
|
|
assert "tasks.knowledge_fs_lifecycle_tasks" in celery_app.conf["imports"]
|
|
assert "tasks.knowledge_fs_source_import_tasks" in celery_app.conf["imports"]
|
|
assert "tasks.delete_conversation_task" in celery_app.conf["imports"]
|
|
assert celery_app.conf["beat_schedule"]["conversation_cleanup_sweeper"] == {
|
|
"task": "tasks.delete_conversation_task.sweep_deleted_conversations",
|
|
"schedule": timedelta(minutes=5),
|
|
}
|
|
assert celery_app.conf["beat_schedule"]["knowledge_fs_staged_upload_cleanup"] == {
|
|
"task": "tasks.knowledge_fs_lifecycle_tasks.cleanup_knowledge_fs_staged_uploads",
|
|
"schedule": timedelta(seconds=2),
|
|
}
|
|
assert celery_app.conf["beat_schedule"]["knowledge_fs_upgrade_file_cleanup"] == {
|
|
"task": "tasks.knowledge_fs_upgrade_tasks.cleanup_deferred_knowledge_fs_upgrade_files",
|
|
"schedule": timedelta(seconds=2),
|
|
}
|
|
assert celery_app.conf["beat_schedule"]["knowledge_fs_upgrade_recovery"] == {
|
|
"task": "tasks.knowledge_fs_upgrade_tasks.recover_stalled_knowledge_fs_upgrades",
|
|
"schedule": timedelta(seconds=60),
|
|
"options": {"expires": 60},
|
|
}
|
|
schedules = celery_app.conf["beat_schedule"]
|
|
if background_enabled:
|
|
assert celery_app.conf["broker_transport_options"]["visibility_timeout"] == 10800
|
|
for operation, queue in OPERATION_QUEUES.items():
|
|
assert schedules[f"knowledge_fs_background_{operation}"] == {
|
|
"task": OPERATION_TASK,
|
|
"schedule": timedelta(seconds=5),
|
|
"kwargs": {"operation": operation},
|
|
"options": {"queue": queue, "expires": 30},
|
|
}
|
|
else:
|
|
assert not any(name.startswith("knowledge_fs_background_") for name in schedules)
|
|
|
|
with (
|
|
patch(
|
|
"services.knowledge_fs.lifecycle_readiness.get_configured_knowledge_fs_lifecycle_worker_readiness",
|
|
return_value=KnowledgeFSLifecycleWorkerReadiness(ready=False, blockers=("test",)),
|
|
),
|
|
):
|
|
preview_only_app = init_app(DifyApp(f"{__name__}.preview_only"))
|
|
|
|
assert "tasks.knowledge_fs_initial_source_preview_tasks" in preview_only_app.conf["imports"]
|
|
assert "tasks.knowledge_fs_failed_retrieval_tasks" in preview_only_app.conf["imports"]
|
|
assert "tasks.knowledge_fs_initial_source_tasks" not in preview_only_app.conf["imports"]
|
|
assert "tasks.knowledge_fs_lifecycle_tasks" not in preview_only_app.conf["imports"]
|
|
assert "tasks.knowledge_fs_source_import_tasks" not in preview_only_app.conf["imports"]
|
|
assert "knowledge_fs_upgrade_recovery" not in preview_only_app.conf["beat_schedule"]
|
|
|
|
@staticmethod
|
|
@pytest.mark.parametrize("visibility", [None, 3600, 7200, 7230])
|
|
def test_long_running_delivery_rejects_an_unsafe_redis_visibility(visibility: int | None) -> None:
|
|
with patch("extensions.ext_celery.dify_config") as config:
|
|
config.CELERY_USE_SENTINEL = False
|
|
config.REDIS_KEY_PREFIX = ""
|
|
config.CELERY_BROKER_URL = "redis://localhost:6379/0"
|
|
config.CELERY_BROKER_VISIBILITY_TIMEOUT = visibility
|
|
config.KNOWLEDGE_FS_BACKGROUND_WORKER_ENABLED = True
|
|
config.KNOWLEDGE_FS_BACKGROUND_TASK_TIMEOUT_SECONDS = 7200
|
|
with pytest.raises(ValueError, match="VISIBILITY_TIMEOUT"):
|
|
get_celery_broker_transport_options()
|
|
|
|
|
|
pytestmark = pytest.mark.usefixtures(preserve_celery_apps.__name__)
|