Add redis tests

antispam/caches/redis/redis.py

@@ -23,7 +23,7 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, List, AsyncIterable
+from typing import TYPE_CHECKING, List, AsyncIterable, Dict
 
 from attr import asdict
 
@@ -69,10 +69,10 @@ async def get_guild(self, guild_id: int) -> Guild:
         # This is actually a dict here
         guild.options = Options(**guild.options)  # type: ignore
 
-        guild_members: List[Member] = []
+        guild_members: Dict[int, Member] = {}
         for member_id in guild.members:  # type: ignore
             member: Member = await self.get_member(member_id, guild_id)
-            guild_members.append(member)
+            guild_members[member.id] = member
 
         guild.members = guild_members
         return guild
@@ -97,7 +97,7 @@ async def get_member(self, member_id: int, guild_id: int) -> Member:
             member_id,
             guild_id,
         )
-        resp = await self.redis.get(f"MEMBER:{member_id}:{guild_id}")
+        resp = await self.redis.get(f"MEMBER:{guild_id}:{member_id}")
         if not resp:
             raise MemberNotFound
 
@@ -109,7 +109,6 @@ async def get_member(self, member_id: int, guild_id: int) -> Member:
             messages.append(Message(**message))  # type: ignore
 
         member.messages = messages
-
         return member
 
     async def set_member(self, member: Member) -> None:
@@ -121,8 +120,12 @@ async def set_member(self, member: Member) -> None:
 
         # Ensure a guild exists
         try:
-            # NOOP
-            await self.get_guild(member.guild_id)
+            guild = await self.get_guild(member.guild_id)
+
+            guild.members = [m.id for m in guild.members.values()]
+            guild.members.append(member.id)
+            guild_as_json = json.dumps(asdict(guild, recurse=True))
+            await self.redis.set(f"GUILD:{guild.id}", guild_as_json)
         except GuildNotFound:
             guild = Guild(id=member.guild_id, options=self.handler.options)
             guild.members = [member.id]
@@ -136,6 +139,13 @@ async def delete_member(self, member_id: int, guild_id: int) -> None:
         log.debug(
             "Attempting to delete Member(id=%s) in Guild(id=%s)", member_id, guild_id
         )
+        try:
+            guild: Guild = await self.get_guild(guild_id)
+            guild.members.pop(member_id)
+            await self.set_guild(guild)
+        except:
+            pass
+
         await self.redis.delete(f"MEMBER:{guild_id}:{member_id}")
 
     async def add_message(self, message: Message) -> None:
@@ -187,6 +197,9 @@ async def get_all_guilds(self) -> AsyncIterable[Guild]:
 
     async def get_all_members(self, guild_id: int) -> AsyncIterable[Member]:
         log.debug("Yielding all cached members for Guild(id=%s)", guild_id)
+        # NOOP
+        await self.get_guild(guild_id)
+
         keys: List[bytes] = await self.redis.keys(f"MEMBER:{guild_id}:*")
         for key in keys:
             key = key.decode("utf-8").split(":")[2]

tests/conftest.py

@@ -8,6 +8,7 @@
 from antispam import AntiSpamHandler, PluginCache, Options
 from antispam.caches import MemoryCache
 from antispam.caches.mongo import MongoCache
+from antispam.caches.redis import RedisCache
 from antispam.core import Core
 from antispam.dataclasses import Guild, Member, Message
 from antispam.enums import Library
@@ -17,7 +18,7 @@
 from examples.custom_multistage_punishments.AntiSpamTrackerSubclass import (
     MyCustomTracker,
 )
-from tests.mocks import MockedMember
+from tests.mocks import MockedMember, MockedRedis
 from tests.mocks.mock_document import MockedDocument
 
 
@@ -126,3 +127,8 @@ def create_mongo_cache(create_handler) -> MockedMongoCache:
 @pytest.fixture()
 def create_example_tracker_subclass(create_handler) -> MyCustomTracker:
     return MyCustomTracker(create_handler, 3, 1000)
+
+
+@pytest.fixture()
+def create_redis_cache(create_handler) -> RedisCache:
+    return RedisCache(create_handler, MockedRedis())

tests/mocks/__init__.py

@@ -3,3 +3,4 @@
 from .mocked_guild import MockedGuild
 from .mocked_member import MockedMember
 from .mocked_message import MockedMessage
+from .mocked_redis import MockedRedis

tests/mocks/mocked_redis.py

@@ -0,0 +1,51 @@
+from typing import Dict, Optional
+
+
+class MockedRedis:
+    """A mock aioredis.Redis class that
+    imitates the required methods.
+    """
+
+    def __init__(self):
+        self._data: Dict[str, dict] = {}
+
+    @property
+    def cache(self) -> Dict:
+        return self._data
+
+    async def get(self, key) -> Optional[dict]:
+        try:
+            return self._data[key]
+        except:
+            return None
+
+    async def set(self, key, value):
+        self._data[key] = value
+
+    async def delete(self, key):
+        try:
+            self._data.pop(key)
+        except:
+            pass
+
+    async def flushdb(self, *args, **kwargs):
+        self._data = {}
+
+    async def keys(self, pattern: str):
+        if pattern.startswith("GUILD"):
+            return self._get_guilds()
+
+        member, guild_id, wildcard = pattern.split(":")
+        return self._get_members(guild_id)
+
+    def _get_guilds(self):
+        return [
+            key.encode("utf-8") for key in self._data.keys() if key.startswith("GUILD")
+        ]
+
+    def _get_members(self, guild_id):
+        return [
+            key.encode("utf-8")
+            for key in self._data.keys()
+            if key.startswith(f"MEMBER:{guild_id}")
+        ]

tests/test_redis_cache.py

@@ -0,0 +1,173 @@
+import orjson as json
+
+import pytest
+from attr import asdict
+
+from antispam import GuildNotFound, MemberNotFound, Options
+from antispam.dataclasses import Guild, Member, Message
+from antispam.enums import ResetType
+from antispam.factory import FactoryBuilder
+
+
+class TestRedisCache:
+    @pytest.mark.asyncio
+    async def test_redis_init(self, create_redis_cache):
+        # Coverage
+        await create_redis_cache.initialize()
+
+    @pytest.mark.asyncio
+    async def test_get_guild(self, create_redis_cache):
+        with pytest.raises(GuildNotFound):
+            await create_redis_cache.get_guild(1)
+
+        create_redis_cache.redis.cache["GUILD:1"] = json.dumps(
+            asdict(Guild(1), recurse=True)
+        )
+
+        val = await create_redis_cache.get_guild(1)
+        assert val
+        assert isinstance(val, Guild)
+
+    @pytest.mark.asyncio
+    async def test_set_guild(self, create_redis_cache):
+        assert not create_redis_cache.redis.cache
+        await create_redis_cache.set_guild(Guild(1, Options()))
+        assert create_redis_cache.redis.cache
+
+    @pytest.mark.asyncio
+    async def test_get_member(self, create_redis_cache):
+        await create_redis_cache.set_guild(Guild(1, Options()))
+        with pytest.raises(MemberNotFound):
+            await create_redis_cache.get_member(1, 1)
+
+        await create_redis_cache.set_member(Member(1, 1))
+
+        val = await create_redis_cache.get_member(1, 1)
+        assert val
+        assert isinstance(val, Member)
+
+    @pytest.mark.asyncio
+    async def test_set_member(self, create_redis_cache):
+        await create_redis_cache.set_member(Member(1, 1))
+        await create_redis_cache.set_member(Member(2, 1))
+
+        r_1 = await create_redis_cache.get_guild(1)
+        assert len(r_1.members) == 2
+
+        await create_redis_cache.set_guild(Guild(1, Options()))
+        r_2 = await create_redis_cache.get_guild(1)
+        assert len(r_2.members) == 0
+
+        await create_redis_cache.set_member(Member(1, 1))
+
+        r_3 = await create_redis_cache.get_guild(1)
+        assert len(r_3.members) == 1
+
+    @pytest.mark.asyncio
+    async def test_add_message_raw(self, create_redis_cache):
+        """Test without a pre-filled cache"""
+        await create_redis_cache.add_message(Message(1, 2, 3, 4, "Content"))
+
+        r_1 = await create_redis_cache.get_guild(3)
+        assert len(r_1.members) == 1
+        assert len(r_1.members[4].messages) == 1
+
+    @pytest.mark.asyncio
+    async def test_add_message_filled_guild(self, create_redis_cache):
+        """Test with an existing guild"""
+        await create_redis_cache.set_guild(Guild(3, Options()))
+        await create_redis_cache.add_message(Message(1, 2, 3, 4, "Content"))
+
+        r_1 = await create_redis_cache.get_member(4, 3)
+        r_2 = await create_redis_cache.get_guild(3)
+        assert len(r_2.members) == 1
+        assert len(r_1.messages) == 1
+
+    @pytest.mark.asyncio
+    async def test_add_message_filled(self, create_redis_cache):
+        """Test with a pre-filled member"""
+        await create_redis_cache.set_member(Member(4, 3))
+        await create_redis_cache.add_message(Message(1, 2, 3, 4, "Content"))
+
+        r_1 = await create_redis_cache.get_guild(3)
+        assert len(r_1.members) == 1
+        assert len(r_1.members[4].messages) == 1
+
+    @pytest.mark.asyncio
+    async def test_reset_member_count(self, create_redis_cache):
+        await create_redis_cache.set_member(Member(1, 1, kick_count=1, warn_count=2))
+
+        member = await create_redis_cache.get_member(1, 1)
+        assert (member.kick_count, member.warn_count) == (1, 2)
+
+        await create_redis_cache.reset_member_count(1, 1, ResetType.KICK_COUNTER)
+
+        member = await create_redis_cache.get_member(1, 1)
+        assert (member.kick_count, member.warn_count) == (0, 2)
+
+        await create_redis_cache.reset_member_count(1, 1, ResetType.WARN_COUNTER)
+
+        member = await create_redis_cache.get_member(1, 1)
+        assert (member.kick_count, member.warn_count) == (0, 0)
+
+    @pytest.mark.asyncio
+    async def test_get_all_members(self, create_redis_cache):
+        with pytest.raises(GuildNotFound):
+            await FactoryBuilder.get_all_members_as_list(create_redis_cache, 1)
+
+        await create_redis_cache.set_member(Member(1, 1))
+        await create_redis_cache.set_member(Member(2, 1))
+        await create_redis_cache.set_member(Member(3, 1))
+
+        members = await FactoryBuilder.get_all_members_as_list(create_redis_cache, 1)
+        assert len(members) == 3
+        assert members == [Member(1, 1), Member(2, 1), Member(3, 1)]
+
+    @pytest.mark.asyncio
+    async def test_get_all_guilds(self, create_redis_cache):
+        guilds = await FactoryBuilder.get_all_guilds_as_list(create_redis_cache)
+        assert len(guilds) == 0
+
+        await create_redis_cache.set_guild(Guild(1, Options()))
+        guilds = await FactoryBuilder.get_all_guilds_as_list(create_redis_cache)
+        assert len(guilds) == 1
+
+        await create_redis_cache.set_guild(Guild(2, Options()))
+        guilds = await FactoryBuilder.get_all_guilds_as_list(create_redis_cache)
+        assert len(guilds) == 2
+        assert guilds == [Guild(1, Options()), Guild(2, Options())]
+
+    @pytest.mark.asyncio
+    async def test_delete_guild(self, create_redis_cache):
+        guilds = await FactoryBuilder.get_all_guilds_as_list(create_redis_cache)
+        assert len(guilds) == 0
+
+        await create_redis_cache.delete_guild(1)
+
+        guilds = await FactoryBuilder.get_all_guilds_as_list(create_redis_cache)
+        assert len(guilds) == 0
+
+        await create_redis_cache.set_guild(Guild(1))
+
+        guilds = await FactoryBuilder.get_all_guilds_as_list(create_redis_cache)
+        assert len(guilds) == 1
+
+        await create_redis_cache.delete_guild(1)
+
+        guilds = await FactoryBuilder.get_all_guilds_as_list(create_redis_cache)
+        assert len(guilds) == 0
+
+    @pytest.mark.asyncio
+    async def test_delete_member(self, create_redis_cache):
+        await create_redis_cache.delete_member(1, 2)
+        await create_redis_cache.set_guild(Guild(2))
+        await create_redis_cache.delete_member(1, 2)
+
+        guild = await create_redis_cache.get_guild(2)
+        guild.members[1] = Member(1, 2)
+        await create_redis_cache.set_guild(guild)
+        assert len(guild.members) == 1
+
+        await create_redis_cache.delete_member(1, 2)
+        g = await create_redis_cache.get_guild(2)
+        assert len(g.members) == 0