dify/api/tests/unit_tests/extensions/test_ext_redis.py

from __future__ import annotations

from collections.abc import Callable, Mapping
from dataclasses import dataclass, field

import pytest
import redis
from redis import RedisError
from redis.cluster import RedisCluster
from redis.retry import Retry

from configs import dify_config
from dify_app import DifyApp
from extensions import ext_redis
from extensions.ext_redis import (
    RedisClientWrapper,
    _get_base_redis_params,
    _get_cluster_connection_health_params,
    _get_connection_health_params,
    _normalize_redis_key_prefix,
    _parse_redis_nodes,
    _serialize_redis_name,
    redis_fallback,
)
from libs.broadcast_channel.redis.pubsub_channel import BroadcastChannel as RedisBroadcastChannel
from libs.broadcast_channel.redis.sharded_channel import ShardedRedisBroadcastChannel
from tests.helpers.unit.config_override import apply_config_overrides


def _initialize_redis(
    monkeypatch: pytest.MonkeyPatch,
    *,
    use_sentinel: bool,
    pubsub_redis_url: str | None,
    pubsub_uses_clusters: bool = False,
    redis_host: str = "default-redis.example",
) -> tuple[DifyApp, redis.Redis[bytes] | RedisCluster[bytes], list[tuple[str, bool]]]:
    """Initialize Redis without network I/O and record dedicated PubSub client creation."""
    monkeypatch.setattr(dify_config, "REDIS_USE_SENTINEL", use_sentinel)
    monkeypatch.setattr(dify_config, "REDIS_USE_CLUSTERS", False)
    monkeypatch.setattr(dify_config, "REDIS_SENTINELS", "sentinel.example:26379")
    monkeypatch.setattr(dify_config, "REDIS_SENTINEL_SERVICE_NAME", "redis-primary")
    monkeypatch.setattr(dify_config, "REDIS_SENTINEL_USERNAME", None)
    monkeypatch.setattr(dify_config, "REDIS_SENTINEL_PASSWORD", None)
    monkeypatch.setattr(dify_config, "REDIS_HOST", redis_host)
    monkeypatch.setattr(dify_config, "REDIS_PORT", 6379)
    monkeypatch.setattr(dify_config, "REDIS_USERNAME", None)
    monkeypatch.setattr(dify_config, "REDIS_PASSWORD", None)
    monkeypatch.setattr(dify_config, "REDIS_DB", 0)
    monkeypatch.setattr(dify_config, "REDIS_USE_SSL", False)
    monkeypatch.setattr(dify_config, "REDIS_ENABLE_CLIENT_SIDE_CACHE", False)
    monkeypatch.setattr(dify_config, "REDIS_MAX_CONNECTIONS", None)
    monkeypatch.setattr(dify_config, "PUBSUB_REDIS_URL", pubsub_redis_url)
    monkeypatch.setattr(dify_config, "PUBSUB_REDIS_USE_CLUSTERS", pubsub_uses_clusters)

    wrapper = ext_redis.RedisClientWrapper()
    monkeypatch.setattr(ext_redis, "redis_client", wrapper)
    monkeypatch.setattr(ext_redis, "_pubsub_redis_client", None)

    pubsub_calls: list[tuple[str, bool]] = []

    def create_pubsub_client(pubsub_url: str, use_clusters: bool) -> redis.Redis[bytes]:
        pubsub_calls.append((pubsub_url, use_clusters))
        return redis.Redis.from_url(pubsub_url)

    monkeypatch.setattr(ext_redis, "_create_pubsub_client", create_pubsub_client)

    app = DifyApp("test-ext-redis")
    ext_redis.init_app(app)

    return app, wrapper._require_client(), pubsub_calls


@pytest.mark.parametrize(
    ("use_sentinel", "configured_url", "expected_dedicated_url"),
    [
        pytest.param(True, None, None, id="sentinel-default"),
        pytest.param(True, "", None, id="sentinel-empty"),
        pytest.param(True, " \t ", None, id="sentinel-whitespace"),
        pytest.param(
            True,
            "  redis://event-bus.example:6380/3  ",
            "redis://event-bus.example:6380/3",
            id="sentinel-explicit",
        ),
        pytest.param(False, None, "redis://default-redis.example:6379/0", id="standalone-default"),
        pytest.param(False, "", "redis://default-redis.example:6379/0", id="standalone-empty"),
        pytest.param(False, " \t ", "redis://default-redis.example:6379/0", id="standalone-whitespace"),
        pytest.param(
            False,
            "  redis://event-bus.example:6380/3  ",
            "redis://event-bus.example:6380/3",
            id="standalone-explicit",
        ),
    ],
)
def test_init_app_selects_pubsub_client_from_redis_mode_and_explicit_url(
    monkeypatch: pytest.MonkeyPatch,
    use_sentinel: bool,
    configured_url: str | None,
    expected_dedicated_url: str | None,
) -> None:
    app, primary_client, pubsub_calls = _initialize_redis(
        monkeypatch,
        use_sentinel=use_sentinel,
        pubsub_redis_url=configured_url,
    )

    assert app.extensions["redis"] is ext_redis.redis_client
    if expected_dedicated_url is None:
        assert pubsub_calls == []
        assert ext_redis._pubsub_redis_client is primary_client
    else:
        assert pubsub_calls == [(expected_dedicated_url, False)]
        assert isinstance(ext_redis._pubsub_redis_client, redis.Redis)
        assert ext_redis._pubsub_redis_client is not primary_client


@pytest.mark.parametrize("configured_url", [None, "", " \t "])
def test_sentinel_without_explicit_pubsub_url_does_not_build_default_url(
    monkeypatch: pytest.MonkeyPatch,
    configured_url: str | None,
) -> None:
    _, primary_client, pubsub_calls = _initialize_redis(
        monkeypatch,
        use_sentinel=True,
        pubsub_redis_url=configured_url,
        redis_host="",
    )

    assert pubsub_calls == []
    assert ext_redis._pubsub_redis_client is primary_client


@pytest.mark.parametrize("use_sentinel", [True, False], ids=["sentinel", "standalone"])
def test_init_app_forwards_pubsub_cluster_setting(
    monkeypatch: pytest.MonkeyPatch,
    use_sentinel: bool,
) -> None:
    _, primary_client, pubsub_calls = _initialize_redis(
        monkeypatch,
        use_sentinel=use_sentinel,
        pubsub_redis_url=" redis://event-bus.example:6380/3 ",
        pubsub_uses_clusters=True,
    )

    assert pubsub_calls == [("redis://event-bus.example:6380/3", True)]
    assert ext_redis._pubsub_redis_client is not primary_client


class TestPubsubChannelCases:
    @staticmethod
    def test_get_pubsub_broadcast_channel_defaults_to_pubsub(monkeypatch: pytest.MonkeyPatch) -> None:
        apply_config_overrides(monkeypatch, PUBSUB_REDIS_CHANNEL_TYPE="pubsub")
        monkeypatch.setattr(ext_redis, "_pubsub_redis_client", object())

        channel = ext_redis.get_pubsub_broadcast_channel()

        assert isinstance(channel, RedisBroadcastChannel)

    @staticmethod
    def test_get_pubsub_broadcast_channel_sharded(monkeypatch: pytest.MonkeyPatch) -> None:
        apply_config_overrides(monkeypatch, PUBSUB_REDIS_CHANNEL_TYPE="sharded")
        monkeypatch.setattr(ext_redis, "_pubsub_redis_client", object())

        channel = ext_redis.get_pubsub_broadcast_channel()

        assert isinstance(channel, ShardedRedisBroadcastChannel)


type _RecordedCall = tuple[str, tuple[object, ...], dict[str, object]]


@dataclass
class _RecordingRedisClient:
    calls: list[_RecordedCall] = field(default_factory=list)

    def _record(self, method: str, *args: object, **kwargs: object) -> None:
        self.calls.append((method, args, kwargs))

    def register_script(self, script: str) -> Callable[..., str]:
        self._record("register_script", script)

        def execute(*, keys: object, args: object, client: object) -> str:
            self._record("script", keys, args, client)
            return "result"

        return execute

    def get(self, name: str | bytes) -> None:
        self._record("get", name)

    def delete(self, *names: str | bytes) -> None:
        self._record("delete", *names)

    def lock(self, name: str, **kwargs: object) -> None:
        self._record("lock", name, **kwargs)

    def hset(self, name: str | bytes, *args: object, **kwargs: object) -> None:
        self._record("hset", name, *args, **kwargs)

    def hgetall(self, name: str | bytes) -> None:
        self._record("hgetall", name)

    def hkeys(self, name: str | bytes) -> None:
        self._record("hkeys", name)

    def hexists(self, name: str | bytes, key: str | bytes) -> None:
        self._record("hexists", name, key)

    def zadd(self, name: str | bytes, mapping: Mapping[object, object], **kwargs: object) -> None:
        self._record("zadd", name, mapping, **kwargs)

    def zrem(self, name: str | bytes, *values: str | bytes) -> None:
        self._record("zrem", name, *values)


class TestRedisCases:
    @staticmethod
    @pytest.fixture(autouse=True)
    def _redis_config(config_overrides: Callable[..., None]) -> None:
        config_overrides(
            REDIS_USERNAME=None,
            REDIS_PASSWORD=None,
            REDIS_DB=0,
            REDIS_SERIALIZATION_PROTOCOL=3,
            REDIS_ENABLE_CLIENT_SIDE_CACHE=False,
            REDIS_RETRY_RETRIES=3,
            REDIS_RETRY_BACKOFF_BASE=1.0,
            REDIS_RETRY_BACKOFF_CAP=10.0,
            REDIS_SOCKET_TIMEOUT=5.0,
            REDIS_SOCKET_CONNECT_TIMEOUT=5.0,
            REDIS_HEALTH_CHECK_INTERVAL=30,
            REDIS_KEEPALIVE=True,
            REDIS_KEEPALIVE_IDLE=60,
            REDIS_KEEPALIVE_INTERVAL=10,
            REDIS_KEEPALIVE_COUNT=3,
            REDIS_KEY_PREFIX="",
        )

    class TestGetConnectionHealthParams:
        def test_includes_all_health_params(self) -> None:
            params = _get_connection_health_params()

            assert "retry" in params
            assert "socket_timeout" in params
            assert "socket_connect_timeout" in params
            assert "health_check_interval" in params
            assert isinstance(params["retry"], Retry)
            assert params["retry"]._retries == 3  # pyrefly: ignore[missing-attribute]
            assert params["socket_timeout"] == 5.0
            assert params["socket_connect_timeout"] == 5.0
            assert params["health_check_interval"] == 30

    class TestGetClusterConnectionHealthParams:
        def test_excludes_health_check_interval(self) -> None:
            params = _get_cluster_connection_health_params()

            assert "retry" in params
            assert "socket_timeout" in params
            assert "socket_connect_timeout" in params
            assert "health_check_interval" not in params

    class TestGetBaseRedisParams:
        def test_includes_retry_and_health_params(self) -> None:
            params = _get_base_redis_params()

            assert "retry" in params
            assert isinstance(params["retry"], Retry)
            assert params["socket_timeout"] == 5.0
            assert params["socket_connect_timeout"] == 5.0
            assert params["health_check_interval"] == 30
            assert params["socket_keepalive"] is True
            assert isinstance(params["socket_keepalive_options"], dict)
            # Existing params still present
            assert params["db"] == 0
            assert params["encoding"] == "utf-8"

    class TestParseRedisNodes:
        def test_trims_nodes(self) -> None:
            assert _parse_redis_nodes("redis-a:6379, redis-b:6380") == [("redis-a", 6379), ("redis-b", 6380)]

        def test_supports_bracketed_ipv6(self) -> None:
            assert _parse_redis_nodes("[2001:db8::10]:6379") == [("2001:db8::10", 6379)]

    class TestRedisFallback:
        def test_redis_fallback_success(self) -> None:
            @redis_fallback(default_return=None)
            def test_func() -> str:
                return "success"

            assert test_func() == "success"

        def test_redis_fallback_error(self) -> None:
            @redis_fallback(default_return="fallback")
            def test_func() -> None:
                raise RedisError("Redis error")

            assert test_func() == "fallback"

        def test_redis_fallback_none_default(self) -> None:
            @redis_fallback()
            def test_func() -> None:
                raise RedisError("Redis error")

            assert test_func() is None

        def test_redis_fallback_with_args(self) -> None:
            received: list[tuple[int, int]] = []

            @redis_fallback(default_return=0)
            def test_func(x: int, y: int) -> None:
                received.append((x, y))
                raise RedisError("Redis error")

            assert test_func(1, 2) == 0
            assert received == [(1, 2)]

        def test_redis_fallback_with_kwargs(self) -> None:
            received: list[tuple[int | None, int | None]] = []

            @redis_fallback(default_return=dict[str, object]())
            def test_func(x: int | None = None, y: int | None = None) -> None:
                received.append((x, y))
                raise RedisError("Redis error")

            assert test_func(x=1, y=2) == {}
            assert received == [(1, 2)]

        def test_redis_fallback_preserves_function_metadata(self) -> None:
            @redis_fallback(default_return=None)
            def test_func() -> None:
                """Test function docstring"""
                pass

            assert test_func.__name__ == "test_func"
            assert test_func.__doc__ == "Test function docstring"

    class TestRedisKeyPrefixHelpers:
        def test_normalize_redis_key_prefix_trims_whitespace(self) -> None:
            assert _normalize_redis_key_prefix("  enterprise-a  ") == "enterprise-a"

        def test_normalize_redis_key_prefix_treats_whitespace_only_as_empty(self) -> None:
            assert _normalize_redis_key_prefix("   ") == ""

        def test_serialize_redis_name_returns_original_when_prefix_empty(self) -> None:
            assert _serialize_redis_name("model_lb_index:test", "") == "model_lb_index:test"

        def test_serialize_redis_name_adds_single_colon_separator(self) -> None:
            assert _serialize_redis_name("model_lb_index:test", "enterprise-a") == "enterprise-a:model_lb_index:test"

    class TestRedisClientWrapperKeyPrefix:
        def test_wrapper_zrem_prefixes_key_but_not_members(self, config_overrides: Callable[..., None]) -> None:
            client = _RecordingRedisClient()
            wrapper = RedisClientWrapper()
            wrapper.initialize(client)  # type: ignore[arg-type]
            config_overrides(REDIS_KEY_PREFIX="enterprise-a")

            wrapper.zrem("reset:user@example.com", "reservation")

            assert client.calls == [("zrem", ("enterprise-a:reset:user@example.com", "reservation"), {})]

        def test_wrapper_registered_script_prefixes_key_arguments(self, config_overrides: Callable[..., None]) -> None:
            raw_client = _RecordingRedisClient()
            wrapper = RedisClientWrapper()
            wrapper.initialize(raw_client)  # type: ignore[arg-type]

            config_overrides(REDIS_KEY_PREFIX="enterprise-a")
            script = wrapper.register_script("return redis.call('GET', KEYS[1])")

            assert script(keys=["device_code:abc"], args=["argument"]) == "result"
            assert raw_client.calls == [
                ("register_script", ("return redis.call('GET', KEYS[1])",), {}),
                ("script", (("enterprise-a:device_code:abc",), ["argument"], raw_client), {}),
            ]

        def test_wrapper_get_prefixes_string_keys(self, config_overrides: Callable[..., None]) -> None:
            client = _RecordingRedisClient()
            wrapper = RedisClientWrapper()
            wrapper.initialize(client)  # type: ignore[arg-type]

            config_overrides(REDIS_KEY_PREFIX="enterprise-a")
            wrapper.get("oauth_state:abc")

            assert client.calls == [("get", ("enterprise-a:oauth_state:abc",), {})]

        def test_wrapper_delete_prefixes_multiple_keys(self, config_overrides: Callable[..., None]) -> None:
            client = _RecordingRedisClient()
            wrapper = RedisClientWrapper()
            wrapper.initialize(client)  # type: ignore[arg-type]

            config_overrides(REDIS_KEY_PREFIX="enterprise-a")
            wrapper.delete("key:a", "key:b")

            assert client.calls == [("delete", ("enterprise-a:key:a", "enterprise-a:key:b"), {})]

        def test_wrapper_lock_prefixes_lock_name(self, config_overrides: Callable[..., None]) -> None:
            client = _RecordingRedisClient()
            wrapper = RedisClientWrapper()
            wrapper.initialize(client)  # type: ignore[arg-type]

            config_overrides(REDIS_KEY_PREFIX="enterprise-a")
            wrapper.lock("resource-lock", timeout=10)

            method, args, kwargs = client.calls[0]
            assert method == "lock"
            assert args == ("enterprise-a:resource-lock",)
            assert kwargs["timeout"] == 10

        def test_wrapper_hash_operations_prefix_key_name(self, config_overrides: Callable[..., None]) -> None:
            client = _RecordingRedisClient()
            wrapper = RedisClientWrapper()
            wrapper.initialize(client)  # type: ignore[arg-type]

            config_overrides(REDIS_KEY_PREFIX="enterprise-a")
            wrapper.hset("hash:key", "field", "value")
            wrapper.hgetall("hash:key")
            wrapper.hkeys("hash:key")
            wrapper.hexists("hash:key", "field")

            assert client.calls == [
                ("hset", ("enterprise-a:hash:key", "field", "value"), {}),
                ("hgetall", ("enterprise-a:hash:key",), {}),
                ("hkeys", ("enterprise-a:hash:key",), {}),
                ("hexists", ("enterprise-a:hash:key", "field"), {}),
            ]

        def test_wrapper_zadd_prefixes_sorted_set_name(self, config_overrides: Callable[..., None]) -> None:
            client = _RecordingRedisClient()
            wrapper = RedisClientWrapper()
            wrapper.initialize(client)  # type: ignore[arg-type]

            config_overrides(REDIS_KEY_PREFIX="enterprise-a")
            wrapper.zadd("zset:key", {"member": 1})

            method, args, kwargs = client.calls[0]
            assert method == "zadd"
            assert args == ("enterprise-a:zset:key", {"member": 1})
            assert kwargs["nx"] is False

        def test_wrapper_preserves_keys_when_prefix_is_empty(self, config_overrides: Callable[..., None]) -> None:
            client = _RecordingRedisClient()
            wrapper = RedisClientWrapper()
            wrapper.initialize(client)  # type: ignore[arg-type]

            config_overrides(REDIS_KEY_PREFIX="   ")
            wrapper.get("plain:key")

            assert client.calls == [("get", ("plain:key",), {})]