mirror of
https://github.com/d3vyce/taskiq-deduplication.git
synced 2026-08-04 19:14:07 +00:00
Version 1.0.0 (#2)
* feat: add taskiq deduplication * doc: rework class comment + update README * fix: make build_deduplication_key private * fix: raise_on_duplicate is now False by default * chore: remove pre_execute * chore: add documentation
This commit is contained in:
d3vyce
@@ -0,0 +1,43 @@
|
||||
import pytest
|
||||
import fakeredis.aioredis
|
||||
from taskiq import TaskiqMessage, TaskiqResult
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def anyio_backend():
|
||||
return "asyncio"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def fake_redis():
|
||||
client = fakeredis.aioredis.FakeRedis()
|
||||
yield client
|
||||
await client.aclose()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def make_message():
|
||||
def _make(task_name="my_task", task_id="task-1", labels=None, kwargs=None):
|
||||
return TaskiqMessage(
|
||||
task_id=task_id,
|
||||
task_name=task_name,
|
||||
labels=labels or {},
|
||||
labels_types={},
|
||||
args=[],
|
||||
kwargs=kwargs or {},
|
||||
)
|
||||
|
||||
return _make
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def make_result():
|
||||
def _make(is_err=False):
|
||||
return TaskiqResult(
|
||||
is_err=is_err,
|
||||
log="",
|
||||
return_value=None,
|
||||
execution_time=0.0,
|
||||
)
|
||||
|
||||
return _make
|
||||
@@ -0,0 +1,242 @@
|
||||
from unittest.mock import AsyncMock, 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.startswith("myapp:locks:")
|
||||
|
||||
|
||||
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}))
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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 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_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()
|
||||
Reference in New Issue
Block a user