AntiMalware Tests - implemented minor feedback

1 files changed, 12 insertions, 12 deletions

diff --git a/tests/bot/cogs/test_antimalware.py b/tests/bot/cogs/test_antimalware.py
index e0aa9d6d2..6e06df0a8 100644
--- a/tests/bot/cogs/test_antimalware.py
+++ b/tests/bot/cogs/test_antimalware.py
@@ -1,4 +1,3 @@
-import logging
 import unittest
 from os.path import splitext
 from unittest.mock import AsyncMock, Mock
@@ -6,7 +5,7 @@ from unittest.mock import AsyncMock, Mock
 from discord import NotFound
 
 from bot.cogs import antimalware
-from bot.constants import AntiMalware as AntiMalwareConfig, Channels, Roles, URLs
+from bot.constants import AntiMalware as AntiMalwareConfig, Channels, STAFF_ROLES, URLs
 from tests.helpers import MockAttachment, MockBot, MockMessage, MockRole
 
 MODULE = "bot.cogs.antimalware"
@@ -31,7 +30,7 @@ class AntiMalwareCogTests(unittest.IsolatedAsyncioTestCase):
 
     async def test_message_without_attachment(self):
         """Messages without attachments should result in no action."""
-        self.assertIsNone(await self.cog.on_message(self.message))
+        await self.cog.on_message(self.message)
         self.message.delete.assert_not_called()
 
     async def test_direct_message_with_attachment(self):
@@ -55,8 +54,8 @@ class AntiMalwareCogTests(unittest.IsolatedAsyncioTestCase):
 
     async def test_message_send_by_staff(self):
         """A message send by a member of staff should be ignored."""
-        moderator_role = MockRole(name="Moderator", id=Roles.moderators)
-        self.message.author.roles.append(moderator_role)
+        staff_role = MockRole(id=STAFF_ROLES[0])
+        self.message.author.roles.append(staff_role)
         attachment = MockAttachment(filename="python.asdfsff")
         self.message.attachments = [attachment]
 
@@ -71,6 +70,7 @@ class AntiMalwareCogTests(unittest.IsolatedAsyncioTestCase):
         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")
 
@@ -107,13 +107,13 @@ class AntiMalwareCogTests(unittest.IsolatedAsyncioTestCase):
         self.message.attachments = [attachment]
         self.message.delete = AsyncMock(side_effect=NotFound(response=Mock(status=""), message=""))
 
-        logger = logging.getLogger(MODULE)
-
-        with self.assertLogs(logger=logger, level="INFO") as logs:
+        with self.assertLogs(logger=antimalware.log, level="INFO") as logs:
             await self.cog.on_message(self.message)
+        self.message.delete.assert_called_once()
         self.assertIn(
             f"INFO:{MODULE}:Tried to delete message `{self.message.id}`, but message could not be found.",
-            logs.output)
+            logs.output
+        )
 
     async def test_message_with_illegal_attachment_logs(self):
         """Deleting a message with an illegal attachment should result in a log."""
@@ -123,9 +123,8 @@ class AntiMalwareCogTests(unittest.IsolatedAsyncioTestCase):
         file_extensions = {splitext(attachment.filename.lower())[1] for attachment in self.message.attachments}
         extensions_blocked = file_extensions - set(AntiMalwareConfig.whitelist)
         blocked_extensions_str = ', '.join(extensions_blocked)
-        logger = logging.getLogger(MODULE)
 
-        with self.assertLogs(logger=logger, level="INFO") as logs:
+        with self.assertLogs(logger=antimalware.log, level="INFO") as logs:
             await self.cog.on_message(self.message)
         self.assertEqual(
             [
@@ -133,7 +132,8 @@ class AntiMalwareCogTests(unittest.IsolatedAsyncioTestCase):
                 f"User '{self.message.author}' ({self.message.author.id}) "
                 f"uploaded blacklisted file(s): {blocked_extensions_str}"
             ],
-            logs.output)
+            logs.output
+        )
 
 
 class AntiMalwareSetupTests(unittest.TestCase):