fix: return types not always returning the correct type

Author: onerandomusername

tests/mocks.py

@@ -57,7 +57,6 @@
     "HashableMixin",
     "CustomMockMixin",
     "ColourMixin",
-    "MockAsyncWebhook",
     "MockAttachment",
     "MockBot",
     "MockCategoryChannel",
@@ -75,6 +74,7 @@
     "MockThread",
     "MockUser",
     "MockVoiceChannel",
+    "MockWebhook",
 ]
 
 
@@ -158,7 +158,7 @@ def generate_mock_message(content=unittest.mock.DEFAULT, *args, **kwargs):
     discord.Thread: lambda *args, **kwargs: MockThread(),
     discord.User: lambda *args, **kwargs: MockUser(),
     discord.VoiceChannel: lambda *args, **kwargs: MockVoiceChannel(),
-    discord.Webhook: lambda *args, **kwargs: MockAsyncWebhook(),
+    discord.Webhook: lambda *args, **kwargs: MockWebhook(),
 }
 
 
@@ -222,10 +222,8 @@ def __init__(self, **kwargs):
             # this list can be added to as methods are discovered
             if attr.__name__ == "send":
                 hints["return"] = discord.Message
-            elif self.__class__ == discord.Message and attr.__name__ == "edit":
-                # set up message editing to return the same object
-                mock_config[f"{attr.__name__}.return_value"] = self
-                continue
+            elif attr.__name__ == "edit":
+                hints["return"] = type(self.spec_set)
 
             if hints.get("return") is None:
                 continue
@@ -822,7 +820,7 @@ def __init__(self, **kwargs) -> None:
 webhook_instance = discord.Webhook(data=unittest.mock.MagicMock(), session=unittest.mock.MagicMock())
 
 
-class MockAsyncWebhook(CustomMockMixin, unittest.mock.NonCallableMagicMock):
+class MockWebhook(CustomMockMixin, unittest.mock.NonCallableMagicMock):
     """
     A MagicMock subclass to mock Webhook objects using an AsyncWebhookAdapter.
 

tests/test_mocks.py

@@ -400,42 +400,56 @@ class TestReturnTypes:
     Eg, ctx.send should return a message object.
     """
 
+    @pytest.mark.parametrize(
+        "mock_cls",
+        [
+            mocks.MockClientUser,
+            mocks.MockGuild,
+            mocks.MockMember,
+            mocks.MockMessage,
+            mocks.MockTextChannel,
+            mocks.MockVoiceChannel,
+            mocks.MockWebhook,
+        ],
+    )
     @pytest.mark.asyncio
-    async def test_message_edit_returns_self(self):
-        """Message editing edits the message in place. We should be returning the message."""
-        msg = mocks.MockMessage()
-
-        new_msg = await msg.edit()
+    async def test_edit_returns_same_class(self, mock_cls):
+        """Edit methods return a new instance of the same type."""
+        mock = mock_cls()
 
-        assert isinstance(new_msg, discord.Message)
+        new_mock = await mock.edit()
 
-        assert msg is new_msg
+        assert isinstance(new_mock, type(mock_cls.spec_set))
 
+    @pytest.mark.parametrize(
+        "mock_cls",
+        [
+            mocks.MockMember,
+            mocks.MockTextChannel,
+            mocks.MockThread,
+            mocks.MockUser,
+        ],
+    )
     @pytest.mark.asyncio
-    async def test_channel_send_returns_message(self):
+    async def test_messageable_send_returns_message(self, mock_cls):
         """Ensure that channel objects return mocked messages when sending messages."""
-        channel = mocks.MockTextChannel()
+        messageable = mock_cls()
 
-        msg = await channel.send("hi")
+        msg = await messageable.send("hi")
 
         print(type(msg))
         assert isinstance(msg, discord.Message)
 
+    @pytest.mark.parametrize(
+        "mock_cls",
+        [mocks.MockMessage, mocks.MockTextChannel],
+    )
     @pytest.mark.asyncio
-    async def test_message_thread_create_returns_thread(self):
-        """Thread create methods should return a MockThread."""
-        msg = mocks.MockMessage()
-
-        thread = await msg.create_thread()
-
-        assert isinstance(thread, discord.Thread)
-
-    @pytest.mark.asyncio
-    async def test_channel_thread_create_returns_thread(self):
+    async def test_thread_create_returns_thread(self, mock_cls):
         """Thread create methods should return a MockThread."""
-        channel = mocks.MockTextChannel()
+        mock = mock_cls()
 
-        thread = await channel.create_thread()
+        thread = await mock.create_thread()
 
         assert isinstance(thread, discord.Thread)