Another attempt to fix pgbouncer error (1093.) (#1097)

* Another attempt to fix pgbouncer error (1093.)

* Fixes for various pgbouncer problems.

* different approach with custom cursor.

* Fix rebase.

* Missed this.

* Fix completion refresher test.

* Black.

* Unused import.

* Changelog.

* Fix race condition in test.

* Switch from is_pgbouncer to more generic is_virtual_database, and duck-type it. Add very dumb unit test for virtual cursor.

* Remove debugger code.

4 files changed, 45 insertions, 12 deletions

diff --git a/tests/features/steps/specials.py b/tests/features/steps/specials.py
index 813292c4..a85f3710 100644
--- a/tests/features/steps/specials.py
+++ b/tests/features/steps/specials.py
@@ -22,5 +22,10 @@ def step_see_refresh_started(context):
     Wait to see refresh output.
     """
     wrappers.expect_pager(
-        context, "Auto-completion refresh started in the background.\r\n", timeout=2
+        context,
+        [
+            "Auto-completion refresh started in the background.\r\n",
+            "Auto-completion refresh restarted.\r\n",
+        ],
+        timeout=2,
     )
diff --git a/tests/features/steps/wrappers.py b/tests/features/steps/wrappers.py
index 78d76881..0ca83669 100644
--- a/tests/features/steps/wrappers.py
+++ b/tests/features/steps/wrappers.py
@@ -39,9 +39,15 @@ def expect_exact(context, expected, timeout):
 
 
 def expect_pager(context, expected, timeout):
+    formatted = expected if isinstance(expected, list) else [expected]
+    formatted = [
+        f"{context.conf['pager_boundary']}\r\n{t}{context.conf['pager_boundary']}\r\n"
+        for t in formatted
+    ]
+
     expect_exact(
         context,
-        "{0}\r\n{1}{0}\r\n".format(context.conf["pager_boundary"], expected),
+        formatted,
         timeout=timeout,
     )
 
diff --git a/tests/test_completion_refresher.py b/tests/test_completion_refresher.py
index 34cf5700..a5529d6b 100644
--- a/tests/test_completion_refresher.py
+++ b/tests/test_completion_refresher.py
@@ -37,7 +37,7 @@ def test_refresh_called_once(refresher):
     :return:
     """
     callbacks = Mock()
-    pgexecute = Mock()
+    pgexecute = Mock(**{"is_virtual_database.return_value": False})
     special = Mock()
 
     with patch.object(refresher, "_bg_refresh") as bg_refresh:
@@ -57,7 +57,7 @@ def test_refresh_called_twice(refresher):
     """
     callbacks = Mock()
 
-    pgexecute = Mock()
+    pgexecute = Mock(**{"is_virtual_database.return_value": False})
     special = Mock()
 
     def dummy_bg_refresh(*args):
@@ -84,14 +84,12 @@ def test_refresh_with_callbacks(refresher):
     :param refresher:
     """
     callbacks = [Mock()]
-    pgexecute_class = Mock()
-    pgexecute = Mock()
+    pgexecute = Mock(**{"is_virtual_database.return_value": False})
     pgexecute.extra_args = {}
     special = Mock()
 
-    with patch("pgcli.completion_refresher.PGExecute", pgexecute_class):
-        # Set refreshers to 0: we're not testing refresh logic here
-        refresher.refreshers = {}
-        refresher.refresh(pgexecute, special, callbacks)
-        time.sleep(1)  # Wait for the thread to work.
-        assert callbacks[0].call_count == 1
+    # Set refreshers to 0: we're not testing refresh logic here
+    refresher.refreshers = {}
+    refresher.refresh(pgexecute, special, callbacks)
+    time.sleep(1)  # Wait for the thread to work.
+    assert callbacks[0].call_count == 1
diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py
index 513e6192..109674cb 100644
--- a/tests/test_pgexecute.py
+++ b/tests/test_pgexecute.py
@@ -520,6 +520,21 @@ class BrokenConnection:
         raise psycopg2.InterfaceError("I'm broken!")
 
 
+class VirtualCursor:
+    """Mock a cursor to virtual database like pgbouncer."""
+
+    def __init__(self):
+        self.protocol_error = False
+        self.protocol_message = ""
+        self.description = None
+        self.status = None
+        self.statusmessage = "Error"
+
+    def execute(self, *args, **kwargs):
+        self.protocol_error = True
+        self.protocol_message = "Command not supported"
+
+
 @dbtest
 def test_exit_without_active_connection(executor):
     quit_handler = MagicMock()
@@ -542,3 +557,12 @@ def test_exit_without_active_connection(executor):
         # an exception should be raised when running a query without active connection
         with pytest.raises(psycopg2.InterfaceError):
             run(executor, "select 1", pgspecial=pgspecial)
+
+
+@dbtest
+def test_virtual_database(executor):
+    virtual_connection = MagicMock()
+    virtual_connection.cursor.return_value = VirtualCursor()
+    with patch.object(executor, "conn", virtual_connection):
+        result = run(executor, "select 1")
+        assert "Command not supported" in result