taskiq-deduplication/tests/test_middleware.py

from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from taskiq_deduplication import DuplicateTaskError, RedisDeduplicationMiddleware
from taskiq_deduplication.middleware import (
    DEDUP_EXPLICIT_KEY_LABEL,
    DEDUP_KEY_FIELDS_LABEL,
    DEDUP_LABEL,
    DEDUP_TTL_LABEL,
)


@pytest.fixture
def middleware(fake_redis):
    mw = RedisDeduplicationMiddleware(redis_url="redis://localhost")
    mw._redis = fake_redis
    return mw


class TestDefaultBuildDeduplicationKey:
    def test_same_kwargs_same_key(self, middleware, make_message):
        m1 = make_message(kwargs={"a": 1, "b": 2})
        m2 = make_message(kwargs={"a": 1, "b": 2})
        assert middleware._build_deduplication_key(
            m1
        ) == middleware._build_deduplication_key(m2)

    def test_different_kwargs_different_key(self, middleware, make_message):
        m1 = make_message(kwargs={"a": 1})
        m2 = make_message(kwargs={"a": 2})
        assert middleware._build_deduplication_key(
            m1
        ) != middleware._build_deduplication_key(m2)

    def test_kwarg_order_invariant(self, middleware, make_message):
        m1 = make_message(kwargs={"a": 1, "b": 2})
        m2 = make_message(kwargs={"b": 2, "a": 1})
        assert middleware._build_deduplication_key(
            m1
        ) == middleware._build_deduplication_key(m2)

    def test_different_task_names_different_keys(self, middleware, make_message):
        m1 = make_message(task_name="task_a", kwargs={"x": 1})
        m2 = make_message(task_name="task_b", kwargs={"x": 1})
        assert middleware._build_deduplication_key(
            m1
        ) != middleware._build_deduplication_key(m2)

    def test_explicit_key_label(self, middleware, make_message):
        m = make_message(labels={DEDUP_EXPLICIT_KEY_LABEL: "my-lock"})
        key = middleware._build_deduplication_key(m)
        assert key == "taskiq:deduplication:my-lock"

    def test_explicit_key_ignores_kwargs(self, middleware, make_message):
        m1 = make_message(kwargs={"a": 1}, labels={DEDUP_EXPLICIT_KEY_LABEL: "fixed"})
        m2 = make_message(kwargs={"a": 99}, labels={DEDUP_EXPLICIT_KEY_LABEL: "fixed"})
        assert middleware._build_deduplication_key(
            m1
        ) == middleware._build_deduplication_key(m2)

    def test_key_fields_filters_kwargs(self, middleware, make_message):
        m1 = make_message(
            kwargs={"a": 1, "b": 2, "c": 3},
            labels={DEDUP_KEY_FIELDS_LABEL: ["a", "b"]},
        )
        m2 = make_message(
            kwargs={"a": 1, "b": 2, "c": 999},
            labels={DEDUP_KEY_FIELDS_LABEL: ["a", "b"]},
        )
        assert middleware._build_deduplication_key(
            m1
        ) == middleware._build_deduplication_key(m2)

    def test_key_fields_different_included_fields(self, middleware, make_message):
        m1 = make_message(
            kwargs={"a": 1, "b": 2},
            labels={DEDUP_KEY_FIELDS_LABEL: ["a"]},
        )
        m2 = make_message(
            kwargs={"a": 1, "b": 99},
            labels={DEDUP_KEY_FIELDS_LABEL: ["a"]},
        )
        assert middleware._build_deduplication_key(
            m1
        ) == middleware._build_deduplication_key(m2)

    def test_key_prefix_in_output(self, make_message):
        mw = RedisDeduplicationMiddleware(
            redis_url="redis://localhost", key_prefix="myapp:locks"
        )
        mw._redis = None
        m = make_message()
        key = mw._build_deduplication_key(m)
        assert key is not None and key.startswith("myapp:locks:")

    def test_empty_kwargs_produces_consistent_key(self, middleware, make_message):
        m1 = make_message(kwargs={})
        m2 = make_message(kwargs={})
        assert middleware._build_deduplication_key(
            m1
        ) == middleware._build_deduplication_key(m2)

    def test_key_fields_empty_list_ignores_all_kwargs(self, middleware, make_message):
        m1 = make_message(kwargs={"a": 1}, labels={DEDUP_KEY_FIELDS_LABEL: []})
        m2 = make_message(kwargs={"a": 999}, labels={DEDUP_KEY_FIELDS_LABEL: []})
        assert middleware._build_deduplication_key(
            m1
        ) == middleware._build_deduplication_key(m2)

    def test_key_fields_absent_from_kwargs_are_ignored(self, middleware, make_message):
        m1 = make_message(
            kwargs={"order_id": 1}, labels={DEDUP_KEY_FIELDS_LABEL: ["user_id"]}
        )
        m2 = make_message(
            kwargs={"order_id": 999}, labels={DEDUP_KEY_FIELDS_LABEL: ["user_id"]}
        )
        assert middleware._build_deduplication_key(
            m1
        ) == middleware._build_deduplication_key(m2)

    def test_explicit_key_takes_precedence_over_key_fields(
        self, middleware, make_message
    ):
        m = make_message(
            kwargs={"a": 1},
            labels={DEDUP_EXPLICIT_KEY_LABEL: "my-lock", DEDUP_KEY_FIELDS_LABEL: ["a"]},
        )
        assert middleware._build_deduplication_key(m) == "taskiq:deduplication:my-lock"

    def test_non_serializable_kwargs_returns_none(self, middleware, make_message):
        m = make_message(kwargs={"dt": object()})
        assert middleware._build_deduplication_key(m) is None


class TestPreSend:
    @pytest.mark.anyio
    async def test_first_send_passes(self, middleware, make_message):
        msg = make_message()
        result = await middleware.pre_send(msg)
        assert result is msg

    @pytest.mark.anyio
    async def test_duplicate_raises(self, middleware, make_message):
        msg = make_message()
        await middleware.pre_send(msg)
        with pytest.raises(DuplicateTaskError):
            await middleware.pre_send(make_message())

    @pytest.mark.anyio
    async def test_deduplication_disabled_label(self, middleware, make_message):
        msg1 = make_message(labels={DEDUP_LABEL: False})
        msg2 = make_message(labels={DEDUP_LABEL: False})
        await middleware.pre_send(msg1)
        await middleware.pre_send(msg2)  # should not raise

    @pytest.mark.anyio
    async def test_deduplication_disabled_by_default_init(
        self, fake_redis, make_message
    ):
        mw = RedisDeduplicationMiddleware(
            redis_url="redis://localhost", default_deduplication=False
        )
        mw._redis = fake_redis
        await mw.pre_send(make_message())
        await mw.pre_send(make_message())  # should not raise

    @pytest.mark.anyio
    async def test_ttl_applied(self, middleware, fake_redis, make_message):
        msg = make_message(labels={DEDUP_TTL_LABEL: 42})
        await middleware.pre_send(msg)
        key = middleware._build_deduplication_key(msg)
        ttl = await fake_redis.ttl(key)
        assert 0 < ttl <= 42

    @pytest.mark.anyio
    async def test_different_kwargs_both_pass(self, middleware, make_message):
        await middleware.pre_send(make_message(kwargs={"x": 1}))
        await middleware.pre_send(make_message(kwargs={"x": 2}))

    @pytest.mark.anyio
    async def test_non_serializable_kwargs_skips_deduplication(
        self, middleware, make_message
    ):
        msg1 = make_message(kwargs={"dt": object()})
        msg2 = make_message(kwargs={"dt": object()})
        await middleware.pre_send(msg1)
        await middleware.pre_send(msg2)  # should not raise

    @pytest.mark.anyio
    async def test_default_ttl_applied(self, fake_redis, make_message):
        mw = RedisDeduplicationMiddleware(redis_url="redis://localhost", default_ttl=77)
        mw._redis = fake_redis
        msg = make_message()
        await mw.pre_send(msg)
        key = mw._build_deduplication_key(msg)
        ttl = await fake_redis.ttl(key)
        assert 0 < ttl <= 77


class TestPostExecute:
    @pytest.mark.anyio
    async def test_releases_lock(
        self, middleware, fake_redis, make_message, make_result
    ):
        msg = make_message()
        await middleware.pre_send(msg)
        key = middleware._build_deduplication_key(msg)
        assert await fake_redis.exists(key)

        await middleware.post_execute(msg, make_result())
        assert not await fake_redis.exists(key)

    @pytest.mark.anyio
    async def test_deduplication_disabled_noop(
        self, middleware, fake_redis, make_message, make_result
    ):
        msg = make_message()
        await middleware.pre_send(msg)
        key = middleware._build_deduplication_key(msg)

        disabled_msg = make_message(labels={DEDUP_LABEL: False})
        await middleware.post_execute(disabled_msg, make_result())
        assert await fake_redis.exists(key)

    @pytest.mark.anyio
    async def test_post_execute_after_ttl_expiry_is_safe(
        self, middleware, fake_redis, make_message, make_result
    ):
        msg = make_message()
        await middleware.pre_send(msg)
        await fake_redis.delete(middleware._build_deduplication_key(msg))
        await middleware.post_execute(msg, make_result())


class TestOnError:
    @pytest.mark.anyio
    async def test_releases_lock_on_error(
        self, middleware, fake_redis, make_message, make_result
    ):
        msg = make_message()
        await middleware.pre_send(msg)
        key = middleware._build_deduplication_key(msg)
        assert await fake_redis.exists(key)

        await middleware.on_error(msg, make_result(is_err=True), RuntimeError("boom"))
        assert not await fake_redis.exists(key)

    @pytest.mark.anyio
    async def test_deduplication_disabled_noop(
        self, middleware, fake_redis, make_message, make_result
    ):
        msg = make_message()
        await middleware.pre_send(msg)
        key = middleware._build_deduplication_key(msg)

        disabled_msg = make_message(labels={DEDUP_LABEL: False})
        await middleware.on_error(
            disabled_msg, make_result(is_err=True), RuntimeError("x")
        )
        assert await fake_redis.exists(key)


class TestRedispatchAfterRelease:
    @pytest.mark.anyio
    async def test_redispatch_after_post_execute(
        self, middleware, make_message, make_result
    ):
        msg = make_message()
        await middleware.pre_send(msg)
        await middleware.post_execute(msg, make_result())
        await middleware.pre_send(make_message())

    @pytest.mark.anyio
    async def test_redispatch_after_on_error(
        self, middleware, make_message, make_result
    ):
        msg = make_message()
        await middleware.pre_send(msg)
        await middleware.on_error(msg, make_result(is_err=True), RuntimeError("boom"))
        await middleware.pre_send(make_message())


class TestAtomicRelease:
    @pytest.mark.anyio
    async def test_only_owner_can_release(self, middleware, fake_redis, make_message):
        owner_msg = make_message(task_id="owner-task")
        key = middleware._build_deduplication_key(owner_msg)

        await fake_redis.set(key, "owner-task", ex=300)

        await middleware._release_if_owned(key, "other-task")
        assert await fake_redis.exists(key)

        await middleware._release_if_owned(key, "owner-task")
        assert not await fake_redis.exists(key)

    @pytest.mark.anyio
    async def test_release_missing_key_is_noop(self, middleware, fake_redis):
        await middleware._release_if_owned(
            "taskiq:deduplication:nonexistent", "some-task"
        )


class TestLifecycle:
    @pytest.mark.anyio
    async def test_startup_creates_redis_client(self):
        mw = RedisDeduplicationMiddleware(redis_url="redis://localhost")
        assert mw._redis is None
        with patch("redis.asyncio.Redis.from_url") as mock_from_url:
            mock_client = AsyncMock()
            mock_client.register_script = MagicMock()
            mock_from_url.return_value = mock_client
            await mw.startup()
            mock_from_url.assert_called_once_with("redis://localhost")
            assert mw._redis is mock_client

    @pytest.mark.anyio
    async def test_shutdown_closes_redis_client(self):
        mw = RedisDeduplicationMiddleware(redis_url="redis://localhost")
        mock_client = AsyncMock()
        mw._redis = mock_client
        await mw.shutdown()
        mock_client.aclose.assert_called_once()

    @pytest.mark.anyio
    async def test_shutdown_without_startup_is_safe(self):
        mw = RedisDeduplicationMiddleware(redis_url="redis://localhost")
        await mw.shutdown()

    @pytest.mark.anyio
    async def test_pre_send_without_startup_raises_runtime_error(self, make_message):
        mw = RedisDeduplicationMiddleware(redis_url="redis://localhost")
        with pytest.raises(RuntimeError, match="startup"):
            await mw.pre_send(make_message())


class TestStartupRetry:
    @pytest.mark.anyio
    async def test_startup_succeeds_after_retries(self):
        mw = RedisDeduplicationMiddleware(
            redis_url="redis://localhost",
            startup_retries=3,
            startup_retry_delay=0.01,
        )
        with patch("redis.asyncio.Redis.from_url") as mock_from_url:
            mock_client = AsyncMock()
            mock_client.ping.side_effect = [
                ConnectionError("fail"),
                ConnectionError("fail"),
                None,
            ]
            mock_client.register_script = MagicMock()
            mock_from_url.return_value = mock_client
            await mw.startup()
            assert mw._redis is mock_client
            assert mock_client.ping.call_count == 3

    @pytest.mark.anyio
    async def test_startup_raises_after_all_retries_exhausted(self):
        mw = RedisDeduplicationMiddleware(
            redis_url="redis://localhost",
            startup_retries=2,
            startup_retry_delay=0.01,
        )
        with patch("redis.asyncio.Redis.from_url") as mock_from_url:
            mock_client = AsyncMock()
            mock_client.ping.side_effect = ConnectionError("refused")
            mock_from_url.return_value = mock_client
            with pytest.raises(ConnectionError, match="2 attempts"):
                await mw.startup()
            assert mock_client.ping.call_count == 2

    @pytest.mark.anyio
    async def test_startup_no_retry_on_first_success(self):
        mw = RedisDeduplicationMiddleware(
            redis_url="redis://localhost",
            startup_retries=3,
            startup_retry_delay=0.01,
        )
        with patch("redis.asyncio.Redis.from_url") as mock_from_url:
            mock_client = AsyncMock()
            mock_client.register_script = MagicMock()
            mock_from_url.return_value = mock_client
            await mw.startup()
            mock_client.ping.assert_called_once()

    @pytest.mark.anyio
    async def test_startup_retry_delay_exponential(self):
        mw = RedisDeduplicationMiddleware(
            redis_url="redis://localhost",
            startup_retries=3,
            startup_retry_delay=0.01,
        )
        with (
            patch("redis.asyncio.Redis.from_url") as mock_from_url,
            patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep,
        ):
            mock_client = AsyncMock()
            mock_client.ping.side_effect = [
                ConnectionError("fail"),
                ConnectionError("fail"),
                None,
            ]
            mock_client.register_script = MagicMock()
            mock_from_url.return_value = mock_client
            await mw.startup()
            assert mock_sleep.call_count == 2
            assert mock_sleep.call_args_list[0].args[0] == 0.01
            assert mock_sleep.call_args_list[1].args[0] == 0.02


class TestLabelTypeParsing:
    @pytest.mark.anyio
    async def test_bool_label_false_disables_dedup(self, fake_redis, make_message):
        mw = RedisDeduplicationMiddleware(redis_url="redis://localhost")
        mw._redis = fake_redis
        await mw.pre_send(make_message(labels={DEDUP_LABEL: False}))
        await mw.pre_send(make_message(labels={DEDUP_LABEL: False}))

    @pytest.mark.anyio
    async def test_bool_label_true_enables_dedup(self, middleware, make_message):
        msg = make_message(labels={DEDUP_LABEL: True})
        await middleware.pre_send(msg)
        with pytest.raises(DuplicateTaskError):
            await middleware.pre_send(make_message(labels={DEDUP_LABEL: True}))

    @pytest.mark.anyio
    async def test_key_fields_list_parsed_correctly(self, middleware, make_message):
        m1 = make_message(
            kwargs={"a": 1, "b": 2, "c": 3},
            labels={DEDUP_KEY_FIELDS_LABEL: ["a", "b"]},
        )
        m2 = make_message(
            kwargs={"a": 1, "b": 2, "c": 999},
            labels={DEDUP_KEY_FIELDS_LABEL: ["a", "b"]},
        )
        assert middleware._build_deduplication_key(
            m1
        ) == middleware._build_deduplication_key(m2)

    @pytest.mark.anyio
    async def test_key_fields_non_list_ignored(self, middleware, make_message):
        m = make_message(
            kwargs={"a": 1},
            labels={DEDUP_KEY_FIELDS_LABEL: "not-a-list"},
        )
        key = middleware._build_deduplication_key(m)
        assert key is not None


class TestKeyCaching:
    @pytest.mark.anyio
    async def test_key_cached_during_pre_send(self, middleware, make_message):
        from taskiq_deduplication.middleware import _CACHED_KEY_LABEL

        msg = make_message()
        await middleware.pre_send(msg)
        assert _CACHED_KEY_LABEL in msg.labels

    @pytest.mark.anyio
    async def test_post_execute_uses_cached_key(
        self, middleware, fake_redis, make_message, make_result
    ):
        from taskiq_deduplication.middleware import _CACHED_KEY_LABEL

        msg = make_message()
        await middleware.pre_send(msg)
        cached_key = msg.labels.get(_CACHED_KEY_LABEL)
        await middleware.post_execute(msg, make_result())
        assert not await fake_redis.exists(cached_key)

    @pytest.mark.anyio
    async def test_on_error_uses_cached_key(
        self, middleware, fake_redis, make_message, make_result
    ):
        from taskiq_deduplication.middleware import _CACHED_KEY_LABEL

        msg = make_message()
        await middleware.pre_send(msg)
        cached_key = msg.labels.get(_CACHED_KEY_LABEL)
        await middleware.on_error(msg, make_result(is_err=True), RuntimeError("boom"))
        assert not await fake_redis.exists(cached_key)

    @pytest.mark.anyio
    async def test_cached_key_none_when_key_build_fails(self, middleware, make_message):
        from taskiq_deduplication.middleware import _CACHED_KEY_LABEL

        msg = make_message(kwargs={"dt": object()})
        await middleware.pre_send(msg)
        assert msg.labels[_CACHED_KEY_LABEL] is None

    @pytest.mark.anyio
    async def test_post_execute_noop_when_cached_key_is_none(
        self, middleware, fake_redis, make_message, make_result
    ):
        msg = make_message(kwargs={"dt": object()})
        await middleware.pre_send(msg)
        await middleware.post_execute(msg, make_result())

    @pytest.mark.anyio
    async def test_on_error_noop_when_cached_key_is_none(
        self, middleware, fake_redis, make_message, make_result
    ):
        msg = make_message(kwargs={"dt": object()})
        await middleware.pre_send(msg)
        await middleware.on_error(msg, make_result(is_err=True), RuntimeError("boom"))

    @pytest.mark.anyio
    async def test_release_if_owned_raises_without_redis(
        self, middleware, make_message
    ):
        middleware._redis = None
        with pytest.raises(RuntimeError, match="startup"):
            await middleware._release_if_owned("some-key", "some-task")