From 33139457b6d254227e400bebb2d6eab83bbded8c Mon Sep 17 00:00:00 2001
From: wangweimin <wang0.618@qq.com>
Date: Sun, 22 May 2022 11:27:25 +0800
Subject: [PATCH] fix: auto hold don't work on script mode

---
 pywebio/platform/path_deploy.py |  1 +
 pywebio/platform/tornado.py     |  8 +++++++-
 test/2.script_mode.py           | 12 ++++++++----
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/pywebio/platform/path_deploy.py b/pywebio/platform/path_deploy.py
index 3509d144..a63aa492 100644
--- a/pywebio/platform/path_deploy.py
+++ b/pywebio/platform/path_deploy.py
@@ -19,6 +19,7 @@
 
 LOCAL_STATIC_URL = '/_pywebio_static'
 
+
 def filename_ok(f):
     return not f.startswith(('.', '_'))
 
diff --git a/pywebio/platform/tornado.py b/pywebio/platform/tornado.py
index a9427fcb..9d785dc3 100644
--- a/pywebio/platform/tornado.py
+++ b/pywebio/platform/tornado.py
@@ -397,6 +397,7 @@ def start_server_in_current_thread_session():
     class SingleSessionWSHandler(_webio_handler(cdn=False)):
         session = None
         instance = None
+        closed = False
 
         def open(self):
             self.main_session = False
@@ -426,6 +427,7 @@ def open(self):
         def on_close(self):
             if SingleSessionWSHandler.session is not None and self.main_session:
                 self.session.close()
+                self.closed = True
                 logger.debug('ScriptModeSession closed')
 
     async def wait_to_stop_loop(server):
@@ -439,7 +441,11 @@ async def wait_to_stop_loop(server):
             alive_none_daemonic_thread_cnt = sum(
                 1 for t in threading.enumerate() if t.is_alive() and not t.isDaemon()
             )
-            await asyncio.sleep(1)
+            await asyncio.sleep(0.5)
+
+        if SingleSessionWSHandler.instance.session.need_keep_alive():
+            while not SingleSessionWSHandler.instance.closed:
+                await asyncio.sleep(0.5)
 
         # 关闭Websocket连接
         # Close the Websocket connection
diff --git a/test/2.script_mode.py b/test/2.script_mode.py
index 18bb5be3..60bc6d09 100644
--- a/test/2.script_mode.py
+++ b/test/2.script_mode.py
@@ -26,12 +26,16 @@ def test(server_proc: subprocess.Popen, browser: Chrome):
 
     template.test_input(browser)
 
-    # script mode 下，此时 server 应停止
-    server_proc.wait(timeout=8)
-
     time.sleep(1)
     template.save_output(browser, '2.script_mode.html',
-                         process_func=lambda i: i.replace('::1', '127.0.0.1'))  # because tornado default bind ipv4 and ipv6 in script mode
+                         # because tornado default bind ipv4 and ipv6 in script mode
+                         process_func=lambda i: i.replace('::1', '127.0.0.1'))
+
+    browser.close()
+    time.sleep(3)
+
+    # script mode 下，此时 server 应停止
+    server_proc.wait(timeout=8)
 
 
 if __name__ == '__main__':