1 files changed, 32 insertions, 0 deletions
diff --git a/catapult/common/py_utils/py_utils/py_utils_unittest.py b/catapult/common/py_utils/py_utils/py_utils_unittest.py
index e045bcab..e614a454 100644
--- a/catapult/common/py_utils/py_utils/py_utils_unittest.py
+++ b/catapult/common/py_utils/py_utils/py_utils_unittest.py
@@ -21,3 +21,35 @@ class PathTest(unittest.TestCase):
 
 def _GetFileInTestDir(file_name):
   return os.path.join(os.path.dirname(__file__), 'test_data', file_name)
+
+
+class WaitForTest(unittest.TestCase):
+
+  def testWaitForTrue(self):
+    def ReturnTrue():
+      return True
+    self.assertTrue(py_utils.WaitFor(ReturnTrue, .1))
+
+  def testWaitForFalse(self):
+    def ReturnFalse():
+      return False
+
+    with self.assertRaises(py_utils.TimeoutException):
+      py_utils.WaitFor(ReturnFalse, .1)
+
+  def testWaitForEventuallyTrue(self):
+    # Use list to pass to inner function in order to allow modifying the
+    # variable from the outer scope.
+    c = [0]
+    def ReturnCounterBasedValue():
+      c[0] += 1
+      return c[0] > 2
+
+    self.assertTrue(py_utils.WaitFor(ReturnCounterBasedValue, .5))
+
+  def testWaitForTrueLambda(self):
+    self.assertTrue(py_utils.WaitFor(lambda: True, .1))
+
+  def testWaitForFalseLambda(self):
+    with self.assertRaises(py_utils.TimeoutException):
+      py_utils.WaitFor(lambda: False, .1)