Adjust attachments rule to use MockMessage, restructure test cases

1 files changed, 23 insertions, 20 deletions

diff --git a/tests/bot/rules/test_attachments.py b/tests/bot/rules/test_attachments.py
index 4bb0acf7c..2f8294922 100644
--- a/tests/bot/rules/test_attachments.py
+++ b/tests/bot/rules/test_attachments.py
@@ -1,26 +1,17 @@
 import asyncio
 import unittest
-from dataclasses import dataclass
-from typing import Any, List
 
 from bot.rules import attachments
+from tests.helpers import MockMessage
 
 
-# Using `MagicMock` sadly doesn't work for this usecase
-# since it's __eq__ compares the MagicMock's ID. We just
-# want to compare the actual attributes we set.
-@dataclass
-class FakeMessage:
-    author: str
-    attachments: List[Any]
-
-
-def msg(total_attachments: int) -> FakeMessage:
-    return FakeMessage(author='lemon', attachments=list(range(total_attachments)))
+def msg(total_attachments: int) -> MockMessage:
+    """Builds a message with `total_attachments` attachments."""
+    return MockMessage(author='lemon', attachments=list(range(total_attachments)))
 
 
 class AttachmentRuleTests(unittest.TestCase):
-    """Tests applying the `attachment` antispam rule."""
+    """Tests applying the `attachments` antispam rule."""
 
     def test_allows_messages_without_too_many_attachments(self):
         """Messages without too many attachments are allowed as-is."""
@@ -38,13 +29,25 @@ class AttachmentRuleTests(unittest.TestCase):
     def test_disallows_messages_with_too_many_attachments(self):
         """Messages with too many attachments trigger the rule."""
         cases = (
-            ((msg(4), msg(0), msg(6)), [msg(4), msg(6)], 10),
-            ((msg(6),), [msg(6)], 6),
-            ((msg(1),) * 6, [msg(1)] * 6, 6),
+            ([msg(4), msg(0), msg(6)], 10),
+            ([msg(6)], 6),
+            ([msg(1)] * 6, 6),
         )
-        for messages, relevant_messages, total in cases:
-            with self.subTest(messages=messages, relevant_messages=relevant_messages, total=total):
-                last_message, *recent_messages = messages
+        for messages, total in cases:
+            last_message, *recent_messages = messages
+            relevant_messages = [last_message] + [
+                msg
+                for msg in recent_messages
+                if msg.author == last_message.author
+                and len(msg.attachments) > 0
+            ]
+
+            with self.subTest(
+                last_message=last_message,
+                recent_messages=recent_messages,
+                relevant_messages=relevant_messages,
+                total=total
+            ):
                 coro = attachments.apply(last_message, recent_messages, {'max': 5})
                 self.assertEqual(
                     asyncio.run(coro),