Update tests/prompts/test_common.py

tmbo · Jul 24, 2024 · c856b3b · c856b3b
1 parent 6fd05b4
commit c856b3b
Showing 1 changed file with 5 additions and 8 deletions.
diff --git a/tests/prompts/test_common.py b/tests/prompts/test_common.py
@@ -87,15 +87,12 @@ async def run(inp):
             == 1000000000000000000000000000001
         )
 
+    pipe_input = create_pipe_input()
     if prompt_toolkit_version < (3, 0, 29):
-        inp = create_pipe_input()
-        try:
-            return asyncio.run(run(inp))
-        finally:
-            inp.close()
-    else:
-        with create_pipe_input() as inp:
-            asyncio.run(run(inp))
+        pipe_input = pipe_input.create()
+
+    with pipe_input as inp:
+        asyncio.run(run(inp))
 
 
 def test_prompt_highlight_coexist():