diff options
Diffstat (limited to 'tests')
| -rw-r--r-- | tests/bot/exts/filters/test_antimalware.py | 45 |
1 files changed, 30 insertions, 15 deletions
diff --git a/tests/bot/exts/filters/test_antimalware.py b/tests/bot/exts/filters/test_antimalware.py index 3393c6cdc..06d78de9d 100644 --- a/tests/bot/exts/filters/test_antimalware.py +++ b/tests/bot/exts/filters/test_antimalware.py @@ -104,24 +104,39 @@ class AntiMalwareCogTests(unittest.IsolatedAsyncioTestCase): 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) + """A message containing a .txt/.json/.csv file should result in the correct embed.""" + test_values = ( + ("text", ".txt"), + ("json", ".json"), + ("csv", ".csv"), + ) - 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) + for file_name, disallowed_extension in test_values: + with self.subTest(file_name=file_name, disallowed_extension=disallowed_extension): + + attachment = MockAttachment(filename=f"{file_name}{disallowed_extension}") + 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( + blocked_extension=disallowed_extension, + cmd_channel_mention=cmd_channel.mention + ) async def test_other_disallowed_extension_embed_description(self): - """Test the description for a non .py/.txt disallowed extension.""" + """Test the description for a non .py/.txt/.json/.csv disallowed extension.""" attachment = MockAttachment(filename="python.disallowed") self.message.attachments = [attachment] self.message.channel.send = AsyncMock() |