tests/extmod: Add test for ThreadSafeFlag.

Signed-off-by: Jim Mussared <jim.mussared@gmail.com>

2 files changed, 100 insertions, 0 deletions

diff --git a/tests/extmod/uasyncio_threadsafeflag.py b/tests/extmod/uasyncio_threadsafeflag.py
new file mode 100644
index 0000000000..4e002a3d2a
--- /dev/null
+++ b/tests/extmod/uasyncio_threadsafeflag.py
@@ -0,0 +1,79 @@
+# Test Event class
+
+try:
+    import uasyncio as asyncio
+except ImportError:
+    print("SKIP")
+    raise SystemExit
+
+
+import micropython
+
+try:
+    micropython.schedule
+except AttributeError:
+    print("SKIP")
+    raise SystemExit
+
+
+try:
+    # Unix port can't select/poll on user-defined types.
+    import uselect as select
+
+    poller = select.poll()
+    poller.register(asyncio.ThreadSafeFlag())
+except TypeError:
+    print("SKIP")
+    raise SystemExit
+
+
+async def task(id, flag):
+    print("task", id)
+    await flag.wait()
+    print("task", id, "done")
+
+
+def set_from_schedule(flag):
+    print("schedule")
+    flag.set()
+    print("schedule done")
+
+
+async def main():
+    flag = asyncio.ThreadSafeFlag()
+
+    # Set the flag from within the loop.
+    t = asyncio.create_task(task(1, flag))
+    print("yield")
+    await asyncio.sleep(0)
+    print("set event")
+    flag.set()
+    print("yield")
+    await asyncio.sleep(0)
+    print("wait task")
+    await t
+
+    # Set the flag from scheduler context.
+    print("----")
+    t = asyncio.create_task(task(2, flag))
+    print("yield")
+    await asyncio.sleep(0)
+    print("set event")
+    micropython.schedule(set_from_schedule, flag)
+    print("yield")
+    await asyncio.sleep(0)
+    print("wait task")
+    await t
+
+    # Flag already set.
+    print("----")
+    print("set event")
+    flag.set()
+    t = asyncio.create_task(task(3, flag))
+    print("yield")
+    await asyncio.sleep(0)
+    print("wait task")
+    await t
+
+
+asyncio.run(main())
diff --git a/tests/extmod/uasyncio_threadsafeflag.py.exp b/tests/extmod/uasyncio_threadsafeflag.py.exp
new file mode 100644
index 0000000000..aef4e479ba
--- /dev/null
+++ b/tests/extmod/uasyncio_threadsafeflag.py.exp
@@ -0,0 +1,21 @@
+yield
+task 1
+set event
+yield
+wait task
+task 1 done
+----
+yield
+task 2
+set event
+yield
+schedule
+schedule done
+wait task
+task 2 done
+----
+set event
+yield
+task 3
+task 3 done
+wait task