Add `return_value` support and assertions to AsyncIteratorMock

The AsyncIteratorMock included in Python 3.8 will work similarly to
the mocks of callabes. This means that it allows you to set the items
it will yield using the `return_value` attribute. It will also have
support for the common Mock-specific assertions.

This commit introduces some backports of those features in a slightly
simplified way to make the transition to Python 3.8 easier in the
future.

1 files changed, 54 insertions, 4 deletions

diff --git a/tests/helpers.py b/tests/helpers.py
index 3e43679fe..50652ef9a 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -127,14 +127,20 @@ class AsyncMock(CustomMockMixin, unittest.mock.MagicMock):
 
 class AsyncIteratorMock:
     """
-    A class to mock asyncronous iterators.
+    A class to mock asynchronous iterators.
 
     This allows async for, which is used in certain Discord.py objects. For example,
-    an async iterator is returned by the Reaction.users() coroutine.
+    an async iterator is returned by the Reaction.users() method.
     """
 
-    def __init__(self, sequence):
-        self.iter = iter(sequence)
+    def __init__(self, iterable: Iterable = None):
+        if iterable is None:
+            iterable = []
+
+        self.iter = iter(iterable)
+        self.iterable = iterable
+
+        self.call_count = 0
 
     def __aiter__(self):
         return self
@@ -145,6 +151,50 @@ class AsyncIteratorMock:
         except StopIteration:
             raise StopAsyncIteration
 
+    def __call__(self):
+        """
+        Keeps track of the number of times an instance has been called.
+
+        This is useful, since it typically shows that the iterator has actually been used somewhere after we have
+        instantiated the mock for an attribute that normally returns an iterator when called.
+        """
+        self.call_count += 1
+        return self
+
+    @property
+    def return_value(self):
+        """Makes `self.iterable` accessible as self.return_value."""
+        return self.iterable
+
+    @return_value.setter
+    def return_value(self, iterable):
+        """Stores the `return_value` as `self.iterable` and its iterator as `self.iter`."""
+        self.iter = iter(iterable)
+        self.iterable = iterable
+
+    def assert_called(self):
+        """Asserts if the AsyncIteratorMock instance has been called at least once."""
+        if self.call_count == 0:
+            raise AssertionError("Expected AsyncIteratorMock to have been called.")
+
+    def assert_called_once(self):
+        """Asserts if the AsyncIteratorMock instance has been called exactly once."""
+        if self.call_count != 1:
+            raise AssertionError(
+                f"Expected AsyncIteratorMock to have been called once. Called {self.call_count} times."
+            )
+
+    def assert_not_called(self):
+        """Asserts if the AsyncIteratorMock instance has not been called."""
+        if self.call_count != 0:
+            raise AssertionError(
+                f"Expected AsyncIteratorMock to not have been called once. Called {self.call_count} times."
+            )
+
+    def reset_mock(self):
+        """Resets the call count, but not the return value or iterator."""
+        self.call_count = 0
+
 
 # Create a guild instance to get a realistic Mock of `discord.Guild`
 guild_data = {