From f3bc041a9d9ce9bbb26304dc251ca51de5487e48 Mon Sep 17 00:00:00 2001
From: Charles Blackmon-Luca <20627856+charlesbluca@users.noreply.github.com>
Date: Wed, 6 Apr 2022 12:41:33 -0700
Subject: [PATCH 1/2] Add handling for newer prompt-toolkit version

---
 tests/integration/test_cmd.py | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/tests/integration/test_cmd.py b/tests/integration/test_cmd.py
index 145b4962f..0f90da41d 100644
--- a/tests/integration/test_cmd.py
+++ b/tests/integration/test_cmd.py
@@ -1,6 +1,8 @@
+import prompt_toolkit
 import pytest
 from dask import config as dask_config
 from mock import MagicMock, patch
+from packaging.version import parse as parseVersion
 from prompt_toolkit.application import create_app_session
 from prompt_toolkit.input import create_pipe_input
 from prompt_toolkit.output import DummyOutput
@@ -8,15 +10,23 @@
 
 from dask_sql.cmd import _meta_commands
 
+_prompt_toolkit_version = parseVersion(prompt_toolkit.__version__)
+
 
 @pytest.fixture(autouse=True, scope="function")
 def mock_prompt_input():
-    pipe_input = create_pipe_input()
-    try:
-        with create_app_session(input=pipe_input, output=DummyOutput()):
-            yield pipe_input
-    finally:
-        pipe_input.close()
+    # TODO: remove if prompt-toolkit min version gets bumped
+    if _prompt_toolkit_version >= parseVersion("3.0.29"):
+        with create_pipe_input() as pipe_input:
+            with create_app_session(input=pipe_input, output=DummyOutput()):
+                yield pipe_input
+    else:
+        pipe_input = create_pipe_input()
+        try:
+            with create_app_session(input=pipe_input, output=DummyOutput()):
+                yield pipe_input
+        finally:
+            pipe_input.close()
 
 
 def _feed_cli_with_input(

From f38a231042cd559da6201823be9f14c134a6ca05 Mon Sep 17 00:00:00 2001
From: Charles Blackmon-Luca <20627856+charlesbluca@users.noreply.github.com>
Date: Wed, 6 Apr 2022 13:48:28 -0700
Subject: [PATCH 2/2] Place compatibility code in _compat

---
 dask_sql/_compat.py           | 6 ++++++
 tests/integration/test_cmd.py | 7 ++-----
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/dask_sql/_compat.py b/dask_sql/_compat.py
index 126416391..ee6d88c2f 100644
--- a/dask_sql/_compat.py
+++ b/dask_sql/_compat.py
@@ -1,6 +1,12 @@
 import pandas as pd
+import prompt_toolkit
 from packaging.version import parse as parseVersion
 
 _pandas_version = parseVersion(pd.__version__)
+_prompt_toolkit_version = parseVersion(prompt_toolkit.__version__)
+
 FLOAT_NAN_IMPLEMENTED = _pandas_version >= parseVersion("1.2.0")
 INT_NAN_IMPLEMENTED = _pandas_version >= parseVersion("1.0.0")
+
+# TODO: remove if prompt-toolkit min version gets bumped
+PIPE_INPUT_CONTEXT_MANAGER = _prompt_toolkit_version >= parseVersion("3.0.29")
diff --git a/tests/integration/test_cmd.py b/tests/integration/test_cmd.py
index 0f90da41d..fdef31bd3 100644
--- a/tests/integration/test_cmd.py
+++ b/tests/integration/test_cmd.py
@@ -1,22 +1,19 @@
-import prompt_toolkit
 import pytest
 from dask import config as dask_config
 from mock import MagicMock, patch
-from packaging.version import parse as parseVersion
 from prompt_toolkit.application import create_app_session
 from prompt_toolkit.input import create_pipe_input
 from prompt_toolkit.output import DummyOutput
 from prompt_toolkit.shortcuts import PromptSession
 
+from dask_sql._compat import PIPE_INPUT_CONTEXT_MANAGER
 from dask_sql.cmd import _meta_commands
 
-_prompt_toolkit_version = parseVersion(prompt_toolkit.__version__)
-
 
 @pytest.fixture(autouse=True, scope="function")
 def mock_prompt_input():
     # TODO: remove if prompt-toolkit min version gets bumped
-    if _prompt_toolkit_version >= parseVersion("3.0.29"):
+    if PIPE_INPUT_CONTEXT_MANAGER:
         with create_pipe_input() as pipe_input:
             with create_app_session(input=pipe_input, output=DummyOutput()):
                 yield pipe_input