mirror of
https://github.com/langgenius/dify.git
synced 2026-02-27 04:57:08 +08:00
2db638b992
Currently, the celery worker executing workflows / chatflows uses redis pubsub to publish events to api. (See \_topic\_msg\_generator and \_publish\_streaming\_response) The current implementation uses the default redis client. For large scale deployment, we need to use a dedicated redis cluster to ensure performance. To achieve this, you should: 1. introduce a dedicated configuration class to control the redis address used for pubsub. (Ideally, there should only be one configuration item such as `pubsub_redis_url`, and its default value should be the original redis confugration.) 2. Add an option to switch between pubsub and sharded pubsub. When shared pubsub is specified, the ShardedRedisBroadcastChannel should be used instead. COmplete the task above, add some unit tests.
110 lines
3.1 KiB
Python
110 lines
3.1 KiB
Python
import os
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from flask import Flask
|
|
|
|
# Getting the absolute path of the current file's directory
|
|
ABS_PATH = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
# Getting the absolute path of the project's root directory
|
|
PROJECT_DIR = os.path.abspath(os.path.join(ABS_PATH, os.pardir, os.pardir))
|
|
|
|
CACHED_APP = Flask(__name__)
|
|
|
|
# set global mock for Redis client
|
|
redis_mock = MagicMock()
|
|
redis_mock.get = MagicMock(return_value=None)
|
|
redis_mock.setex = MagicMock()
|
|
redis_mock.setnx = MagicMock()
|
|
redis_mock.delete = MagicMock()
|
|
redis_mock.lock = MagicMock()
|
|
redis_mock.exists = MagicMock(return_value=False)
|
|
redis_mock.set = MagicMock()
|
|
redis_mock.expire = MagicMock()
|
|
redis_mock.hgetall = MagicMock(return_value={})
|
|
redis_mock.hdel = MagicMock()
|
|
redis_mock.incr = MagicMock(return_value=1)
|
|
|
|
# Ensure OpenDAL fs writes to tmp to avoid polluting workspace
|
|
os.environ.setdefault("OPENDAL_SCHEME", "fs")
|
|
os.environ.setdefault("OPENDAL_FS_ROOT", "/tmp/dify-storage")
|
|
os.environ.setdefault("STORAGE_TYPE", "opendal")
|
|
|
|
# Add the API directory to Python path to ensure proper imports
|
|
import sys
|
|
|
|
sys.path.insert(0, PROJECT_DIR)
|
|
|
|
from extensions import ext_redis
|
|
|
|
|
|
def _patch_redis_clients_on_loaded_modules():
|
|
"""Ensure any module-level redis_client references point to the shared redis_mock."""
|
|
|
|
import sys
|
|
|
|
for module in list(sys.modules.values()):
|
|
if module is None:
|
|
continue
|
|
if hasattr(module, "redis_client"):
|
|
module.redis_client = redis_mock
|
|
if hasattr(module, "pubsub_redis_client"):
|
|
module.pubsub_redis_client = redis_mock
|
|
|
|
|
|
@pytest.fixture
|
|
def app() -> Flask:
|
|
return CACHED_APP
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _provide_app_context(app: Flask):
|
|
with app.app_context():
|
|
yield
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _patch_redis_clients():
|
|
"""Patch redis_client to MagicMock only for unit test executions."""
|
|
|
|
with (
|
|
patch.object(ext_redis, "redis_client", redis_mock),
|
|
patch.object(ext_redis, "pubsub_redis_client", redis_mock),
|
|
):
|
|
_patch_redis_clients_on_loaded_modules()
|
|
yield
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_redis_mock():
|
|
"""reset the Redis mock before each test"""
|
|
redis_mock.reset_mock()
|
|
redis_mock.get.return_value = None
|
|
redis_mock.setex.return_value = None
|
|
redis_mock.setnx.return_value = None
|
|
redis_mock.delete.return_value = None
|
|
redis_mock.exists.return_value = False
|
|
redis_mock.set.return_value = None
|
|
redis_mock.expire.return_value = None
|
|
redis_mock.hgetall.return_value = {}
|
|
redis_mock.hdel.return_value = None
|
|
redis_mock.incr.return_value = 1
|
|
|
|
# Keep any imported modules pointing at the mock between tests
|
|
_patch_redis_clients_on_loaded_modules()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_secret_key():
|
|
"""Ensure SECRET_KEY-dependent logic sees an empty config value by default."""
|
|
|
|
from configs import dify_config
|
|
|
|
original = dify_config.SECRET_KEY
|
|
dify_config.SECRET_KEY = ""
|
|
try:
|
|
yield
|
|
finally:
|
|
dify_config.SECRET_KEY = original
|