+class AssertWithTimeout(object):
+ """Allow some time for an assertion to pass."""
+
+ def __init__(self, timeout=0):
+ self.timeout = timeout
+
+ def __iter__(self):
+ self.deadline = time.time() + self.timeout
+ self.done = False
+ return self
+
+ def next(self):
+ if self.done:
+ raise StopIteration
+ return self.attempt
+
+ def attempt(self, fn, *args, **kwargs):
+ try:
+ fn(*args, **kwargs)
+ except AssertionError:
+ if time.time() > self.deadline:
+ raise
+ time.sleep(0.1)
+ else:
+ self.done = True
+
+