langchain/libs/community/tests/integration_tests/storage/test_redis.py

"""Implement integration tests for Redis storage."""
import os
import typing
import uuid

import pytest

from langchain_community.storage.redis import RedisStore

if typing.TYPE_CHECKING:
    from redis import Redis

pytest.importorskip("redis")


@pytest.fixture
def redis_client() -> Redis:
    """Yield redis client."""
    import redis

    # Using standard port, but protecting against accidental data loss
    # by requiring a password.
    # This fixture flushes the database!
    # The only role of the password is to prevent users from accidentally
    # deleting their data.
    # The password should establish the identity of the server being.
    port = 6379
    password = os.environ.get("REDIS_PASSWORD") or str(uuid.uuid4())
    client = redis.Redis(host="localhost", port=port, password=password, db=0)
    try:
        client.ping()
    except redis.exceptions.ConnectionError:
        pytest.skip(
            "Redis server is not running or is not accessible. "
            "Verify that credentials are correct. "
        )
    # ATTENTION: This will delete all keys in the database!
    client.flushdb()
    return client


def test_mget(redis_client: Redis) -> None:
    """Test mget method."""
    store = RedisStore(client=redis_client, ttl=None)
    keys = ["key1", "key2"]
    redis_client.mset({"key1": b"value1", "key2": b"value2"})
    result = store.mget(keys)
    assert result == [b"value1", b"value2"]


def test_mset(redis_client: Redis) -> None:
    """Test that multiple keys can be set."""
    store = RedisStore(client=redis_client, ttl=None)
    key_value_pairs = [("key1", b"value1"), ("key2", b"value2")]
    store.mset(key_value_pairs)
    result = redis_client.mget(["key1", "key2"])
    assert result == [b"value1", b"value2"]


def test_mdelete(redis_client: Redis) -> None:
    """Test that deletion works as expected."""
    store = RedisStore(client=redis_client, ttl=None)
    keys = ["key1", "key2"]
    redis_client.mset({"key1": b"value1", "key2": b"value2"})
    store.mdelete(keys)
    result = redis_client.mget(keys)
    assert result == [None, None]


def test_yield_keys(redis_client: Redis) -> None:
    store = RedisStore(client=redis_client, ttl=None)
    redis_client.mset({"key1": b"value1", "key2": b"value2"})
    assert sorted(store.yield_keys()) == ["key1", "key2"]
    assert sorted(store.yield_keys(prefix="key*")) == ["key1", "key2"]
    assert sorted(store.yield_keys(prefix="lang*")) == []


def test_namespace(redis_client: Redis) -> None:
    """Test that a namespace is prepended to all keys properly."""
    store = RedisStore(client=redis_client, ttl=None, namespace="meow")
    key_value_pairs = [("key1", b"value1"), ("key2", b"value2")]
    store.mset(key_value_pairs)

    assert sorted(redis_client.scan_iter("*")) == [
        b"meow/key1",
        b"meow/key2",
    ]

    store.mdelete(["key1"])

    assert sorted(redis_client.scan_iter("*")) == [
        b"meow/key2",
    ]

    assert list(store.yield_keys()) == ["key2"]
    assert list(store.yield_keys(prefix="key*")) == ["key2"]
    assert list(store.yield_keys(prefix="key1")) == []