Python: test Awaitable, framework for async test

python/ql/test/experimental/dataflow/coverage/classes.py

@@ -1253,17 +1253,17 @@ def test_exit():
   with With_exit():
     pass
 
-# # 3.4.1. Awaitable Objects
-# # object.__await__(self)
-# class With_await:
+# 3.4.1. Awaitable Objects
+# object.__await__(self)
+class With_await:
 
-#   def __await__(self):
-#     OK()
-#     return "" # edit to match type
+  def __await__(self):
+    OK()
+    return (yield from asyncio.coroutine(lambda: "")())
 
-# async def test_await():
-#   with_await = With_await()
-#   await(with_await) # edit to effect call
+async def atest_await():
+  with_await = With_await()
+  await(with_await)
 
 
 # # 3.4.2. Coroutine Objects

python/ql/test/experimental/dataflow/coverage/validTest.py

@@ -15,6 +15,18 @@ def check_test_function(f):
     sys.stdout = old_stdout
     check_output(capturer.getvalue(), f)
 
+def check_async_test_function(f):
+    from io import StringIO
+    import sys
+    import asyncio
+
+    capturer = StringIO()
+    old_stdout = sys.stdout
+    sys.stdout = capturer
+    asyncio.run(f())
+    sys.stdout = old_stdout
+    check_output(capturer.getvalue(), f)
+
 def check_classes_valid(testFile):
     # import python.ql.test.experimental.dataflow.coverage.classes as tests
     # import classes as tests
@@ -28,6 +40,12 @@ def check_classes_valid(testFile):
                 print("Checking", testFile, item)
                 check_test_function(item)
 
+        elif i.startswith("atest_"):
+            item = getattr(tests,i)
+            if callable(item):
+                print("Checking", testFile, item)
+                check_async_test_function(item)
+
 if __name__ == '__main__':
     check_classes_valid("classes")
     check_classes_valid("test")