1 files changed, 90 insertions, 7 deletions

diff --git a/tests/features/steps/basic_commands.py b/tests/features/steps/basic_commands.py
index 7c87814..687bdc0 100644
--- a/tests/features/steps/basic_commands.py
+++ b/tests/features/steps/basic_commands.py
@@ -64,13 +64,83 @@ def step_ctrl_d(context):
     """
     Send Ctrl + D to hopefully exit.
     """
+    step_try_to_ctrl_d(context)
+    context.cli.expect(pexpect.EOF, timeout=5)
+    context.exit_sent = True
+
+
+@when('we try to send "ctrl + d"')
+def step_try_to_ctrl_d(context):
+    """
+    Send Ctrl + D, perhaps exiting, perhaps not (if a transaction is
+    ongoing).
+    """
     # turn off pager before exiting
     context.cli.sendcontrol("c")
     context.cli.sendline(r"\pset pager off")
     wrappers.wait_prompt(context)
     context.cli.sendcontrol("d")
-    context.cli.expect(pexpect.EOF, timeout=5)
-    context.exit_sent = True
+
+
+@when('we send "ctrl + c"')
+def step_ctrl_c(context):
+    """Send Ctrl + c to hopefully interrupt."""
+    context.cli.sendcontrol("c")
+
+
+@then("we see cancelled query warning")
+def step_see_cancelled_query_warning(context):
+    """
+    Make sure we receive the warning that the current query was cancelled.
+    """
+    wrappers.expect_exact(context, "cancelled query", timeout=2)
+
+
+@then("we see ongoing transaction message")
+def step_see_ongoing_transaction_error(context):
+    """
+    Make sure we receive the warning that a transaction is ongoing.
+    """
+    context.cli.expect("A transaction is ongoing.", timeout=2)
+
+
+@when("we send sleep query")
+def step_send_sleep_15_seconds(context):
+    """
+    Send query to sleep for 15 seconds.
+    """
+    context.cli.sendline("select pg_sleep(15)")
+
+
+@when("we check for any non-idle sleep queries")
+def step_check_for_active_sleep_queries(context):
+    """
+    Send query to check for any non-idle pg_sleep queries.
+    """
+    context.cli.sendline(
+        "select state from pg_stat_activity where query not like '%pg_stat_activity%' and query like '%pg_sleep%' and state != 'idle';"
+    )
+
+
+@then("we don't see any non-idle sleep queries")
+def step_no_active_sleep_queries(context):
+    """Confirm that any pg_sleep queries are either idle or not active."""
+    wrappers.expect_exact(
+        context,
+        context.conf["pager_boundary"]
+        + "\r"
+        + dedent(
+            """
+            +-------+\r
+            | state |\r
+            |-------|\r
+            +-------+\r
+            SELECT 0\r
+        """
+        )
+        + context.conf["pager_boundary"],
+        timeout=5,
+    )
 
 
 @when(r'we send "\?" command')
@@ -131,18 +201,31 @@ def step_see_found(context):
     )
 
 
-@then("we confirm the destructive warning")
-def step_confirm_destructive_command(context):
-    """Confirm destructive command."""
+@then("we respond to the destructive warning: {response}")
+def step_resppond_to_destructive_command(context, response):
+    """Respond to destructive command."""
     wrappers.expect_exact(
         context,
-        "You're about to run a destructive command.\r\nDo you want to proceed? (y/n):",
+        "You're about to run a destructive command.\r\nDo you want to proceed? [y/N]:",
         timeout=2,
     )
-    context.cli.sendline("y")
+    context.cli.sendline(response.strip())
 
 
 @then("we send password")
 def step_send_password(context):
     wrappers.expect_exact(context, "Password for", timeout=5)
     context.cli.sendline(context.conf["pass"] or "DOES NOT MATTER")
+
+
+@when('we send "{text}"')
+def step_send_text(context, text):
+    context.cli.sendline(text)
+    # Try to detect whether we are exiting. If so, set `exit_sent`
+    # so that `after_scenario` correctly cleans up.
+    try:
+        context.cli.expect(pexpect.EOF, timeout=0.2)
+    except pexpect.TIMEOUT:
+        pass
+    else:
+        context.exit_sent = True