allow switching multiple times

Author: maxfischer2781

docs/source/devel/testloop.rst

@@ -27,6 +27,12 @@ Async commands
 
 .. autoclass:: Schedule(*await Any)
 
-.. autoclass:: Switch
+.. py:class::  Switch(skip: int, /)
+    :no-index:
+
+.. py:class::  Switch(min: int, max: int, /)
+    :no-index:
+
+.. autoclass:: Switch()
 
 .. autoclass:: Lock

unittests/test_itertools.py

@@ -322,8 +322,7 @@ async def test_tee_concurrent_locked():
     async def iter_values():
         for item in items:
             # switch to other tasks a few times to guarantees another runs
-            for _ in range(5):
-                await Switch()
+            await Switch(5)
             yield item
 
     async def test_peer(peer_tee):
@@ -354,8 +353,7 @@ async def test_tee_concurrent_unlocked():
     async def iter_values():
         for item in items:
             # switch to other tasks a few times to guarantee another runs
-            for _ in range(5):
-                await Switch()
+            await Switch(5)
             yield item
 
     async def test_peer(peer_tee):

unittests/utility.py

@@ -12,6 +12,7 @@
 )
 from functools import wraps
 from collections import deque
+from random import randint
 
 
 T = TypeVar("T")
@@ -116,10 +117,23 @@ class Switch:
     all other runnable coroutines of the event loop.
     This is similar to the common ``sleep(0)`` function
     of regular event loop frameworks.
+
+    If a single argument is given, this specifies how many
+    turns should be skipped. The default corresponds to `0`.
+    If two arguments are given, this is interpreted as an
+    inclusive interval to randomly select the skip count.
     """
 
+    def __init__(self, skip: int = 0, limit: int = 0, /) -> None:
+        if limit <= 0:
+            self._idle_count = skip
+        else:
+            self._idle_count = randint(skip, limit)
+
     def __await__(self):
         yield self
+        for _ in range(self._idle_count):
+            yield self
 
 
 class Lock: