Repository: python-discord/bot

Issue:
Write unit tests for `bot/cogs/antimalware.py`
Write unit tests for [`bot/cogs/antimalware.py`](../blob/master/bot/cogs/antimalware.py).

## Implementation details
Please make sure to read the general information in the [meta issue](553) and the [testing README](../blob/master/tests/README.md). We are aiming for a 100% [branch coverage](https://coverage.readthedocs.io/en/stable/branch.html) for this file, but if you think that is not possible, please discuss that in this issue.

## Additional information
If you want to work on this issue, **please make sure that you get assigned to it** by one of the core devs before starting to work on it. We would like to prevent the situation that multiple people are working on the same issue. To get assigned, leave a comment showing your interesting in tackling this issue.


Discussion:

Resolving patch:
diff --git a/bot/cogs/antimalware.py b/bot/cogs/antimalware.py
--- a/bot/cogs/antimalware.py
+++ b/bot/cogs/antimalware.py
@@ -1,4 +1,5 @@
 import logging
+import typing as t
 from os.path import splitext
 
 from discord import Embed, Message, NotFound
@@ -9,6 +10,27 @@
 
 log = logging.getLogger(__name__)
 
+PY_EMBED_DESCRIPTION = (
+    "It looks like you tried to attach a Python file - "
+    f"please use a code-pasting service such as {URLs.site_schema}{URLs.site_paste}"
+)
+
+TXT_EMBED_DESCRIPTION = (
+    "**Uh-oh!** It looks like your message got zapped by our spam filter. "
+    "We currently don't allow `.txt` attachments, so here are some tips to help you travel safely: \n\n"
+    "• If you attempted to send a message longer than 2000 characters, try shortening your message "
+    "to fit within the character limit or use a pasting service (see below) \n\n"
+    "• If you tried to show someone your code, you can use codeblocks \n(run `!code-blocks` in "
+    "{cmd_channel_mention} for more information) or use a pasting service like: "
+    f"\n\n{URLs.site_schema}{URLs.site_paste}"
+)
+
+DISALLOWED_EMBED_DESCRIPTION = (
+    "It looks like you tried to attach file type(s) that we do not allow ({blocked_extensions_str}). "
+    f"We currently allow the following file types: **{', '.join(AntiMalwareConfig.whitelist)}**.\n\n"
+    "Feel free to ask in {meta_channel_mention} if you think this is a mistake."
+)
+
 
 class AntiMalware(Cog):
     """Delete messages which contain attachments with non-whitelisted file extensions."""
@@ -29,34 +51,20 @@ async def on_message(self, message: Message) -> None:
             return
 
         embed = Embed()
-        file_extensions = {splitext(attachment.filename.lower())[1] for attachment in message.attachments}
-        extensions_blocked = file_extensions - set(AntiMalwareConfig.whitelist)
+        extensions_blocked = self.get_disallowed_extensions(message)
         blocked_extensions_str = ', '.join(extensions_blocked)
         if ".py" in extensions_blocked:
             # Short-circuit on *.py files to provide a pastebin link
-            embed.description = (
-                "It looks like you tried to attach a Python file - "
-                f"please use a code-pasting service such as {URLs.site_schema}{URLs.site_paste}"
-            )
+            embed.description = PY_EMBED_DESCRIPTION
         elif ".txt" in extensions_blocked:
             # Work around Discord AutoConversion of messages longer than 2000 chars to .txt
             cmd_channel = self.bot.get_channel(Channels.bot_commands)
-            embed.description = (
-                "**Uh-oh!** It looks like your message got zapped by our spam filter. "
-                "We currently don't allow `.txt` attachments, so here are some tips to help you travel safely: \n\n"
-                "• If you attempted to send a message longer than 2000 characters, try shortening your message "
-                "to fit within the character limit or use a pasting service (see below) \n\n"
-                "• If you tried to show someone your code, you can use codeblocks \n(run `!code-blocks` in "
-                f"{cmd_channel.mention} for more information) or use a pasting service like: "
-                f"\n\n{URLs.site_schema}{URLs.site_paste}"
-            )
+            embed.description = TXT_EMBED_DESCRIPTION.format(cmd_channel_mention=cmd_channel.mention)
         elif extensions_blocked:
-            whitelisted_types = ', '.join(AntiMalwareConfig.whitelist)
             meta_channel = self.bot.get_channel(Channels.meta)
-            embed.description = (
-                f"It looks like you tried to attach file type(s) that we do not allow ({blocked_extensions_str}). "
-                f"We currently allow the following file types: **{whitelisted_types}**.\n\n"
-                f"Feel free to ask in {meta_channel.mention} if you think this is a mistake."
+            embed.description = DISALLOWED_EMBED_DESCRIPTION.format(
+                blocked_extensions_str=blocked_extensions_str,
+                meta_channel_mention=meta_channel.mention,
             )
 
         if embed.description:
@@ -73,6 +81,13 @@ async def on_message(self, message: Message) -> None:
             except NotFound:
                 log.info(f"Tried to delete message `{message.id}`, but message could not be found.")
 
+    @classmethod
+    def get_disallowed_extensions(cls, message: Message) -> t.Iterable[str]:
+        """Get an iterable containing all the disallowed extensions of attachments."""
+        file_extensions = {splitext(attachment.filename.lower())[1] for attachment in message.attachments}
+        extensions_blocked = file_extensions - set(AntiMalwareConfig.whitelist)
+        return extensions_blocked
+
 
 def setup(bot: Bot) -> None:
     """Load the AntiMalware cog."""


Test patch:
diff --git a/tests/bot/cogs/test_antimalware.py b/tests/bot/cogs/test_antimalware.py
new file mode 100644
--- /dev/null
+++ b/tests/bot/cogs/test_antimalware.py
@@ -0,0 +1,159 @@
+import unittest
+from unittest.mock import AsyncMock, Mock, patch
+
+from discord import NotFound
+
+from bot.cogs import antimalware
+from bot.constants import AntiMalware as AntiMalwareConfig, Channels, STAFF_ROLES
+from tests.helpers import MockAttachment, MockBot, MockMessage, MockRole
+
+MODULE = "bot.cogs.antimalware"
+
+
+@patch(f"{MODULE}.AntiMalwareConfig.whitelist", new=[".first", ".second", ".third"])
+class AntiMalwareCogTests(unittest.IsolatedAsyncioTestCase):
+    """Test the AntiMalware cog."""
+
+    def setUp(self):
+        """Sets up fresh objects for each test."""
+        self.bot = MockBot()
+        self.cog = antimalware.AntiMalware(self.bot)
+        self.message = MockMessage()
+
+    async def test_message_with_allowed_attachment(self):
+        """Messages with allowed extensions should not be deleted"""
+        attachment = MockAttachment(filename=f"python{AntiMalwareConfig.whitelist[0]}")
+        self.message.attachments = [attachment]
+
+        await self.cog.on_message(self.message)
+        self.message.delete.assert_not_called()
+
+    async def test_message_without_attachment(self):
+        """Messages without attachments should result in no action."""
+        await self.cog.on_message(self.message)
+        self.message.delete.assert_not_called()
+
+    async def test_direct_message_with_attachment(self):
+        """Direct messages should have no action taken."""
+        attachment = MockAttachment(filename="python.disallowed")
+        self.message.attachments = [attachment]
+        self.message.guild = None
+
+        await self.cog.on_message(self.message)
+
+        self.message.delete.assert_not_called()
+
+    async def test_message_with_illegal_extension_gets_deleted(self):
+        """A message containing an illegal extension should send an embed."""
+        attachment = MockAttachment(filename="python.disallowed")
+        self.message.attachments = [attachment]
+
+        await self.cog.on_message(self.message)
+
+        self.message.delete.assert_called_once()
+
+    async def test_message_send_by_staff(self):
+        """A message send by a member of staff should be ignored."""
+        staff_role = MockRole(id=STAFF_ROLES[0])
+        self.message.author.roles.append(staff_role)
+        attachment = MockAttachment(filename="python.disallowed")
+        self.message.attachments = [attachment]
+
+        await self.cog.on_message(self.message)
+
+        self.message.delete.assert_not_called()
+
+    async def test_python_file_redirect_embed_description(self):
+        """A message containing a .py file should result in an embed redirecting the user to our paste site"""
+        attachment = MockAttachment(filename="python.py")
+        self.message.attachments = [attachment]
+        self.message.channel.send = AsyncMock()
+
+        await self.cog.on_message(self.message)
+        self.message.channel.send.assert_called_once()
+        args, kwargs = self.message.channel.send.call_args
+        embed = kwargs.pop("embed")
+
+        self.assertEqual(embed.description, antimalware.PY_EMBED_DESCRIPTION)
+
+    async def test_txt_file_redirect_embed_description(self):
+        """A message containing a .txt file should result in the correct embed."""
+        attachment = MockAttachment(filename="python.txt")
+        self.message.attachments = [attachment]
+        self.message.channel.send = AsyncMock()
+        antimalware.TXT_EMBED_DESCRIPTION = Mock()
+        antimalware.TXT_EMBED_DESCRIPTION.format.return_value = "test"
+
+        await self.cog.on_message(self.message)
+        self.message.channel.send.assert_called_once()
+        args, kwargs = self.message.channel.send.call_args
+        embed = kwargs.pop("embed")
+        cmd_channel = self.bot.get_channel(Channels.bot_commands)
+
+        self.assertEqual(embed.description, antimalware.TXT_EMBED_DESCRIPTION.format.return_value)
+        antimalware.TXT_EMBED_DESCRIPTION.format.assert_called_with(cmd_channel_mention=cmd_channel.mention)
+
+    async def test_other_disallowed_extention_embed_description(self):
+        """Test the description for a non .py/.txt disallowed extension."""
+        attachment = MockAttachment(filename="python.disallowed")
+        self.message.attachments = [attachment]
+        self.message.channel.send = AsyncMock()
+        antimalware.DISALLOWED_EMBED_DESCRIPTION = Mock()
+        antimalware.DISALLOWED_EMBED_DESCRIPTION.format.return_value = "test"
+
+        await self.cog.on_message(self.message)
+        self.message.channel.send.assert_called_once()
+        args, kwargs = self.message.channel.send.call_args
+        embed = kwargs.pop("embed")
+        meta_channel = self.bot.get_channel(Channels.meta)
+
+        self.assertEqual(embed.description, antimalware.DISALLOWED_EMBED_DESCRIPTION.format.return_value)
+        antimalware.DISALLOWED_EMBED_DESCRIPTION.format.assert_called_with(
+            blocked_extensions_str=".disallowed",
+            meta_channel_mention=meta_channel.mention
+        )
+
+    async def test_removing_deleted_message_logs(self):
+        """Removing an already deleted message logs the correct message"""
+        attachment = MockAttachment(filename="python.disallowed")
+        self.message.attachments = [attachment]
+        self.message.delete = AsyncMock(side_effect=NotFound(response=Mock(status=""), message=""))
+
+        with self.assertLogs(logger=antimalware.log, level="INFO"):
+            await self.cog.on_message(self.message)
+        self.message.delete.assert_called_once()
+
+    async def test_message_with_illegal_attachment_logs(self):
+        """Deleting a message with an illegal attachment should result in a log."""
+        attachment = MockAttachment(filename="python.disallowed")
+        self.message.attachments = [attachment]
+
+        with self.assertLogs(logger=antimalware.log, level="INFO"):
+            await self.cog.on_message(self.message)
+
+    async def test_get_disallowed_extensions(self):
+        """The return value should include all non-whitelisted extensions."""
+        test_values = (
+            ([], []),
+            (AntiMalwareConfig.whitelist, []),
+            ([".first"], []),
+            ([".first", ".disallowed"], [".disallowed"]),
+            ([".disallowed"], [".disallowed"]),
+            ([".disallowed", ".illegal"], [".disallowed", ".illegal"]),
+        )
+
+        for extensions, expected_disallowed_extensions in test_values:
+            with self.subTest(extensions=extensions, expected_disallowed_extensions=expected_disallowed_extensions):
+                self.message.attachments = [MockAttachment(filename=f"filename{extension}") for extension in extensions]
+                disallowed_extensions = self.cog.get_disallowed_extensions(self.message)
+                self.assertCountEqual(disallowed_extensions, expected_disallowed_extensions)
+
+
+class AntiMalwareSetupTests(unittest.TestCase):
+    """Tests setup of the `AntiMalware` cog."""
+
+    def test_setup(self):
+        """Setup of the extension should call add_cog."""
+        bot = MockBot()
+        antimalware.setup(bot)
+        bot.add_cog.assert_called_once()