rasahq--rasa/tests/core/test_lock_store.py

import asyncio
import logging
import sys
import time
from pathlib import Path
from typing import Text
from unittest.mock import Mock, patch

import numpy as np
import pytest
import rasa.core.lock_store
from _pytest.logging import LogCaptureFixture
from _pytest.monkeypatch import MonkeyPatch
from rasa.core.agent import Agent
from rasa.core.channels import UserMessage
from rasa.core.constants import DEFAULT_LOCK_LIFETIME
from rasa.core.lock import TicketLock
from rasa.core.lock_store import (
    DEFAULT_REDIS_LOCK_STORE_KEY_PREFIX,
    InMemoryLockStore,
    LockError,
    LockStore,
    RedisLockStore,
)
from rasa.shared.constants import INTENT_MESSAGE_PREFIX
from rasa.shared.exceptions import ConnectionException
from rasa.utils.endpoints import EndpointConfig, read_endpoint_config


class FakeRedisLockStore(RedisLockStore):
    """Fake `RedisLockStore` using `fakeredis` library."""

    # skipcq: PYL-W0231
    # noinspection PyMissingConstructor
    def __init__(self):
        import fakeredis

        self.red = fakeredis.FakeStrictRedis()

        # added in redis==3.3.0, but not yet in fakeredis
        self.red.connection_pool.connection_class.health_check_interval = 0

        self.key_prefix = DEFAULT_REDIS_LOCK_STORE_KEY_PREFIX


def test_issue_ticket():
    lock = TicketLock("random id 0")

    # no lock issued
    assert lock.last_issued == -1
    assert lock.now_serving == 0

    # no one is waiting
    assert not lock.is_someone_waiting()

    # issue ticket
    ticket = lock.issue_ticket(1)
    assert ticket == 0
    assert lock.last_issued == 0
    assert lock.now_serving == 0

    # someone is waiting
    assert lock.is_someone_waiting()


def test_remove_expired_tickets():
    lock = TicketLock("random id 1")

    # issue one long- and one short-lived ticket
    _ = list(map(lock.issue_ticket, [k for k in [0.01, 10]]))

    # both tickets are there
    assert len(lock.tickets) == 2

    # sleep and only one ticket should be left
    time.sleep(0.02)
    lock.remove_expired_tickets()
    assert len(lock.tickets) == 1


@pytest.mark.parametrize("lock_store", [InMemoryLockStore(), FakeRedisLockStore()])
def test_create_lock_store(lock_store: LockStore):
    conversation_id = "my id 0"

    # create and lock
    lock = lock_store.create_lock(conversation_id)
    lock_store.save_lock(lock)
    lock = lock_store.get_lock(conversation_id)
    assert lock
    assert lock.conversation_id == conversation_id


def test_raise_connection_exception_redis_lock_store(monkeypatch: MonkeyPatch):
    monkeypatch.setattr(
        rasa.core.lock_store, "RedisLockStore", Mock(side_effect=ConnectionError())
    )

    with pytest.raises(ConnectionException):
        LockStore.create(
            EndpointConfig(username="username", password="password", type="redis")
        )


@pytest.mark.parametrize("lock_store", [InMemoryLockStore(), FakeRedisLockStore()])
def test_serve_ticket(lock_store: LockStore):
    conversation_id = "my id 1"

    lock = lock_store.create_lock(conversation_id)
    lock_store.save_lock(lock)

    # issue ticket with long lifetime
    ticket_0 = lock_store.issue_ticket(conversation_id, 10)
    assert ticket_0 == 0

    lock = lock_store.get_lock(conversation_id)
    assert lock.last_issued == ticket_0
    assert lock.now_serving == ticket_0
    assert lock.is_someone_waiting()

    # issue another ticket
    ticket_1 = lock_store.issue_ticket(conversation_id, 10)

    # finish serving ticket_0
    lock_store.finish_serving(conversation_id, ticket_0)

    lock = lock_store.get_lock(conversation_id)

    assert lock.last_issued == ticket_1
    assert lock.now_serving == ticket_1
    assert lock.is_someone_waiting()

    # serve second ticket and no one should be waiting
    lock_store.finish_serving(conversation_id, ticket_1)

    lock = lock_store.get_lock(conversation_id)
    assert not lock.is_someone_waiting()


# noinspection PyProtectedMember
@pytest.mark.parametrize("lock_store", [InMemoryLockStore(), FakeRedisLockStore()])
def test_lock_expiration(lock_store: LockStore):
    conversation_id = "my id 2"
    lock = lock_store.create_lock(conversation_id)
    lock_store.save_lock(lock)

    # issue ticket with long lifetime
    ticket = lock.issue_ticket(10)
    assert ticket == 0
    assert not lock._ticket_for_ticket_number(ticket).has_expired()

    # issue ticket with short lifetime
    ticket = lock.issue_ticket(0.00001)
    time.sleep(0.00002)
    assert ticket == 1
    assert lock._ticket_for_ticket_number(ticket) is None

    # newly assigned ticket should get number 1 again
    assert lock.issue_ticket(10) == 1


async def test_multiple_conversation_ids(default_agent: Agent):
    text = INTENT_MESSAGE_PREFIX + 'greet{"name":"Rasa"}'

    conversation_ids = [f"conversation {i}" for i in range(2)]

    # ensure conversations are processed in order
    tasks = [default_agent.handle_text(text, sender_id=_id) for _id in conversation_ids]
    results = await asyncio.gather(*tasks)

    assert results
    processed_ids = [result[0]["recipient_id"] for result in results]
    assert processed_ids == conversation_ids


@pytest.mark.xfail(
    sys.platform == "win32",
    reason="This test sometimes fails on Windows. We want to investigate it further",
)
async def test_message_order(tmp_path: Path, default_agent: Agent):
    start_time = time.time()
    n_messages = 10
    lock_wait = 0.5

    # let's write the incoming order of messages and the order of results to temp files
    results_file = tmp_path / "results_file"
    incoming_order_file = tmp_path / "incoming_order_file"

    # We need to mock `Agent.handle_message()` so we can introduce an
    # artificial holdup (`wait_time_in_seconds`). In the mocked method, we'll
    # record messages as they come and and as they're processed in files so we
    # can check the order later on. We don't need the return value of this method so
    # we'll just return None.
    async def mocked_handle_message(self, message: UserMessage, wait: float) -> None:
        # write incoming message to file
        with open(str(incoming_order_file), "a+") as f_0:
            f_0.write(message.text + "\n")

        async with self.lock_store.lock(
            message.sender_id, wait_time_in_seconds=lock_wait
        ):
            # hold up the message processing after the lock has been acquired
            await asyncio.sleep(wait)

            # write message to file as it's processed
            with open(str(results_file), "a+") as f_1:
                f_1.write(message.text + "\n")

            return None

    # We'll send n_messages from the same sender_id with different blocking times
    # after the lock has been acquired.
    # We have to ensure that the messages are processed in the right order.
    with patch.object(Agent, "handle_message", mocked_handle_message):
        # use decreasing wait times so that every message after the first one
        # does not acquire its lock immediately
        wait_times = np.linspace(0.1, 0.05, n_messages)
        tasks = [
            default_agent.handle_message(
                UserMessage(f"sender {i}", sender_id="some id"), wait=k
            )
            for i, k in enumerate(wait_times)
        ]

        # execute futures
        await asyncio.gather(*(asyncio.ensure_future(t) for t in tasks))

        expected_order = [f"sender {i}" for i in range(len(wait_times))]

        # ensure order of incoming messages is as expected
        with open(str(incoming_order_file)) as f:
            incoming_order = [line for line in f.read().split("\n") if line]
            assert incoming_order == expected_order

        # ensure results are processed in expected order
        with open(str(results_file)) as f:
            results_order = [line for line in f.read().split("\n") if line]
            assert results_order == expected_order

        # Every message after the first one will wait `lock_wait` seconds to acquire its
        # lock (`wait_time_in_seconds` kwarg in `lock_store.lock()`).
        # Let's make sure that this is not blocking and test that total test
        # execution time is less than  the sum of all wait times plus
        # (n_messages - 1) * lock_wait
        time_limit = np.sum(wait_times[1:])
        time_limit += (n_messages - 1) * lock_wait
        assert time.time() - start_time < time_limit


@pytest.mark.xfail(
    sys.platform == "win32",
    reason="This test sometimes fails on Windows. We want to investigate it further",
)
async def test_lock_error(default_agent: Agent):
    lock_lifetime = 0.01
    wait_time_in_seconds = 0.01
    holdup = 0.5

    # Mock message handler again to add a wait time holding up the lock
    # after it's been acquired
    async def mocked_handle_message(self, message: UserMessage) -> None:
        async with self.lock_store.lock(
            message.sender_id,
            wait_time_in_seconds=wait_time_in_seconds,
            lock_lifetime=lock_lifetime,
        ):
            # hold up the message processing after the lock has been acquired
            await asyncio.sleep(holdup)

        return None

    with patch.object(Agent, "handle_message", mocked_handle_message):
        # first message blocks the lock for `holdup`,
        # meaning the second message will not be able to acquire a lock
        tasks = [
            default_agent.handle_message(
                UserMessage(f"sender {i}", sender_id="some id")
            )
            for i in range(2)
        ]

        with pytest.raises(LockError):
            await asyncio.gather(*(asyncio.ensure_future(t) for t in tasks))


async def test_lock_lifetime_environment_variable(monkeypatch: MonkeyPatch):
    import rasa.core.lock_store

    # by default lock lifetime is `DEFAULT_LOCK_LIFETIME`
    assert rasa.core.lock_store._get_lock_lifetime() == DEFAULT_LOCK_LIFETIME

    # set new lock lifetime as environment variable
    new_lock_lifetime = 123
    monkeypatch.setenv("TICKET_LOCK_LIFETIME", str(new_lock_lifetime))

    assert rasa.core.lock_store._get_lock_lifetime() == new_lock_lifetime


@pytest.mark.parametrize("lock_store", [InMemoryLockStore(), FakeRedisLockStore()])
async def test_acquire_lock_debug_message(
    lock_store: LockStore, caplog: LogCaptureFixture
):
    conversation_id = "test_acquire_lock_debug_message"
    wait_time_in_seconds = 0.01

    async def locking_task() -> None:
        async with lock_store.lock(
            conversation_id, wait_time_in_seconds=wait_time_in_seconds
        ):
            # Do a very short sleep so that the other tasks can try to acquire the lock
            # in the meantime
            await asyncio.sleep(0.0)

    with caplog.at_level(logging.DEBUG):
        await asyncio.gather(
            locking_task(),  # Gets served immediately
            locking_task(),  # Gets served second
            locking_task(),  # Gets served last
        )

    assert any(
        f"because 1 other item(s) for this conversation ID have to be finished "
        f"processing first. Retrying in {wait_time_in_seconds} seconds ..." in message
        for message in caplog.messages
    )

    assert any(
        f"because 2 other item(s) for this conversation ID have to be finished "
        f"processing first. Retrying in {wait_time_in_seconds} seconds ..." in message
        for message in caplog.messages
    )


async def test_redis_lock_store_timeout(monkeypatch: MonkeyPatch):
    import redis.exceptions

    lock_store = FakeRedisLockStore()
    monkeypatch.setattr(
        lock_store,
        lock_store.get_or_create_lock.__name__,
        Mock(side_effect=redis.exceptions.TimeoutError),
    )

    with pytest.raises(LockError):
        async with lock_store.lock("some sender"):
            pass


async def test_redis_lock_store_with_invalid_prefix(monkeypatch: MonkeyPatch):
    import redis.exceptions

    lock_store = FakeRedisLockStore()

    prefix = "!asdf234 34#"
    lock_store._set_key_prefix(prefix)
    assert lock_store.key_prefix == DEFAULT_REDIS_LOCK_STORE_KEY_PREFIX

    monkeypatch.setattr(
        lock_store,
        lock_store.get_or_create_lock.__name__,
        Mock(side_effect=redis.exceptions.TimeoutError),
    )

    with pytest.raises(LockError):
        async with lock_store.lock("some sender"):
            pass


async def test_redis_lock_store_with_valid_prefix(monkeypatch: MonkeyPatch):
    import redis.exceptions

    lock_store = FakeRedisLockStore()

    prefix = "chatbot42"
    lock_store._set_key_prefix(prefix)
    assert lock_store.key_prefix == prefix + ":" + DEFAULT_REDIS_LOCK_STORE_KEY_PREFIX

    monkeypatch.setattr(
        lock_store,
        lock_store.get_or_create_lock.__name__,
        Mock(side_effect=redis.exceptions.TimeoutError),
    )

    with pytest.raises(LockError):
        async with lock_store.lock("some sender"):
            pass


def test_create_lock_store_from_endpoint_config(endpoints_path: Text):
    store = read_endpoint_config(endpoints_path, endpoint_type="lock_store")
    tracker_store = RedisLockStore(
        host="localhost",
        port=6379,
        db=0,
        username="username",
        password="password",
        use_ssl=True,
        ssl_keyfile="keyfile.key",
        ssl_certfile="certfile.crt",
        ssl_ca_certs="my-bundle.ca-bundle",
        key_prefix="lock",
    )

    assert isinstance(tracker_store, type(LockStore.create(store)))