Make bot.py Matrix backend quit after 5 seconds

1 files changed, 16 insertions, 2 deletions

diff --git a/bot.py b/bot.py
index 86e14f1..0081c06 100644
--- a/bot.py
+++ b/bot.py
@@ -1,4 +1,5 @@
 from argparse import ArgumentParser
+from asyncio import create_task, sleep, run
 from random import randint, choice
 from re import sub
 
@@ -96,6 +97,9 @@ while True:
     post = sub('(@[^ ]*)@[^ ]*', '\\1', post)
 
 
+    if args.yes:
+        quit()
+
     # Prompt the user
     res = input('Post/Retry/New input/Custom input/Quit: ')
     if res not in 'prnPRNcC':
@@ -134,6 +138,16 @@ for backend, instance, token in zip(args.backend, args.instance, args.token):
         @bot.listener.on_startup
         async def room_joined(room_id):
             await bot.api.send_text_message(room_id=room_id, message=post)
-            quit()
 
-        bot.run()
+        async def wait_quit():
+            await sleep(5)
+            quit()
+        
+        async def run_bot():
+            run = create_task(bot.main())
+            wait = create_task(wait_quit())
+
+            await run
+            await wait
+        
+        run(run_bot())