test: move psql_as_superuser as util and reuse it

test/io/test_io.py

@@ -3,7 +3,6 @@
 import os
 import re
 import signal
-import subprocess
 import time
 import pytest
 import requests
@@ -16,6 +15,7 @@ from util import (
     relativeSeconds,
     drain_stdout,
     match_log,
+    psql_as_superuser,
 )
 from postgrest import (
     Admin,
@@ -33,20 +33,6 @@ from postgrest import (
 )
 
 
-def psql_as_superuser(query):
-    subprocess.check_call(
-        [
-            "psql",
-            "--username",
-            "postgres",
-            "--set",
-            "ON_ERROR_STOP=1",
-            "-c",
-            query,
-        ]
-    )
-
-
 def test_connect_with_dburi(dburi, defaultenv):
     "Connecting with db-uri instead of LIPQ* environment variables should work."
     defaultenv_without_libpq = {
@@ -684,24 +670,15 @@ def test_listener_query_is_visible_in_pg_stat_activity(defaultenv):
     }
 
     with run(env=env):
-        query = """
+        output = psql_as_superuser(
-select query
+            """
-from pg_stat_activity
+        select query
-where application_name = 'listener-query-test'
+        from pg_stat_activity
-  and query = 'LISTEN "pgrst"'
+        where application_name = 'listener-query-test'
-limit 1;
+          and query = 'LISTEN "pgrst"'
-"""
+        limit 1;
-        output = subprocess.check_output(
+        """,
-            [
+            capture_output=True,
-                "psql",
-                "--set",
-                "ON_ERROR_STOP=1",
-                "--tuples-only",
-                "--no-align",
-                "-c",
-                query,
-            ],
-            text=True,
         ).strip()
 
         assert output == 'LISTEN "pgrst"'

test/io/util.py

@@ -1,6 +1,7 @@
 import re
 import threading
 import jwt
+import subprocess
 from datetime import datetime, timedelta, timezone
 
 
@@ -75,3 +76,21 @@ def parse_server_timings_header(header):
         _, duration = duration_text.split("=")
         timings[name.strip()] = float(duration)
     return timings
+
+
+def psql_as_superuser(query, capture_output=False):
+    cmd = [
+        "psql",
+        "--username",
+        "postgres",
+        "--set",
+        "ON_ERROR_STOP=1",
+    ]
+    if capture_output:
+        cmd.extend(["--tuples-only", "--no-align"])
+    cmd.extend(["-c", query])
+
+    if capture_output:
+        return subprocess.check_output(cmd, text=True)
+
+    subprocess.check_call(cmd)