diff --git a/test_runner/batch_others/test_proxy.py b/test_runner/batch_others/test_proxy.py index d5b267273b..ad13a1dde5 100644 --- a/test_runner/batch_others/test_proxy.py +++ b/test_runner/batch_others/test_proxy.py @@ -1,4 +1,54 @@ +import pytest +import subprocess +import signal +import time def test_proxy_select_1(static_proxy): static_proxy.safe_psql("select 1;") + + +def test_proxy_cancel(static_proxy): + """Test that we can cancel a big generate_series query.""" + conn = static_proxy.connect() + conn.cancel() + + with conn.cursor() as cur: + from psycopg2.errors import QueryCanceled + with pytest.raises(QueryCanceled): + cur.execute("select * from generate_series(1, 100000000);") + + +def test_proxy_pgbench_cancel(static_proxy, pg_bin): + """Test that we can cancel the init phase of pgbench.""" + start_time = static_proxy.safe_psql("select now();")[0] + + def get_running_queries(): + magic_string = "fsdsdfhdfhfgbcbfgbfgbf" + with static_proxy.connect() as conn: + with conn.cursor() as cur: + cur.execute(f""" + -- {magic_string} + select query + from pg_stat_activity + where pg_stat_activity.query_start > %s + """, start_time) + return [ + row[0] + for row in cur.fetchall() + if not magic_string in row[0] + ] + + # Let pgbench init run for 1 second + p = subprocess.Popen(['pgbench', '-s500', '-i', static_proxy.connstr()]) + time.sleep(1) + + # Make sure something is still running, and that get_running_queries works + assert len(get_running_queries()) > 0 + + # Send sigint, which would cancel any pgbench queries + p.send_signal(signal.SIGINT) + + # Assert that nothing is running + time.sleep(1) + assert len(get_running_queries()) == 0