Merge pull request #237 from SeitaBV/integer-repeat-option

feat: integer repeat

AUTHORS.rst

@@ -10,4 +10,4 @@ Development Lead
 Contributors
 ------------
 
-None yet. Why not be the first?
+* Felix Claessen <[email protected]>

README.rst

@@ -140,7 +140,11 @@ for convenience use *payload* argument to mock out json response. Example below.
 
 **Repeat response for the same url**  
 
-E.g. for cases you want to test retrying mechanisms
+E.g. for cases where you want to test retrying mechanisms.
+
+- By default, ``repeat=False`` means the response is not repeated (``repeat=1`` does the same).
+- Use ``repeat=n`` to repeat a response n times.
+- Use ``repeat=True`` to repeat a response indefinitely.
 
 .. code:: python
 
@@ -152,14 +156,16 @@ E.g. for cases you want to test retrying mechanisms
     def test_multiple_responses(m):
         loop = asyncio.get_event_loop()
         session = aiohttp.ClientSession()
-        m.get('http://example.com', status=500, repeat=True)
-        m.get('http://example.com', status=200)  # will not take effect
+        m.get('http://example.com', status=500, repeat=2)
+        m.get('http://example.com', status=200)  # will take effect after two preceding calls
 
         resp1 = loop.run_until_complete(session.get('http://example.com'))
         resp2 = loop.run_until_complete(session.get('http://example.com'))
+        resp3 = loop.run_until_complete(session.get('http://example.com'))
 
         assert resp1.status == 500
         assert resp2.status == 500
+        assert resp3.status == 200
 
 
 **match URLs with regular expressions**

aioresponses/core.py

@@ -76,7 +76,7 @@ def __init__(self, url: Union[URL, str, Pattern],
                  content_type: str = 'application/json',
                  response_class: Optional[Type[ClientResponse]] = None,
                  timeout: bool = False,
-                 repeat: bool = False,
+                 repeat: Union[bool, int] = False,
                  reason: Optional[str] = None,
                  callback: Optional[Callable] = None):
         if isinstance(url, Pattern):
@@ -309,7 +309,7 @@ def add(self, url: 'Union[URL, str, Pattern]', method: str = hdrs.METH_GET,
             payload: Optional[Dict] = None,
             headers: Optional[Dict] = None,
             response_class: Optional[Type[ClientResponse]] = None,
-            repeat: bool = False,
+            repeat: Union[bool, int] = False,
             timeout: bool = False,
             reason: Optional[str] = None,
             callback: Optional[Callable] = None) -> None:
@@ -465,8 +465,13 @@ async def match(
             else:
                 return None
 
-            if matcher.repeat is False:
-                del self._matches[key]
+            if isinstance(matcher.repeat, bool):
+                if not matcher.repeat:
+                    del self._matches[key]
+            else:
+                if matcher.repeat == 1:
+                    del self._matches[key]
+                matcher.repeat -= 1
 
             if self.is_exception(response_or_exc):
                 raise response_or_exc

tests/test_aioresponses.py

@@ -610,6 +610,23 @@ async def test_assert_called_twice(self, m: aioresponses):
         with self.assertRaises(AssertionError):
             m.assert_called_once()
 
+    @aioresponses()
+    async def test_integer_repeat_once(self, m: aioresponses):
+        m.get(self.url, repeat=1)
+        m.assert_not_called()
+        await self.session.get(self.url)
+        with self.assertRaises(ClientConnectionError):
+            await self.session.get(self.url)
+
+    @aioresponses()
+    async def test_integer_repeat_twice(self, m: aioresponses):
+        m.get(self.url, repeat=2)
+        m.assert_not_called()
+        await self.session.get(self.url)
+        await self.session.get(self.url)
+        with self.assertRaises(ClientConnectionError):
+            await self.session.get(self.url)
+
     @aioresponses()
     async def test_assert_any_call(self, m: aioresponses):
         http_bin_url = "http://httpbin.org"