Finish making PG tests use a dynamic server address (#3277)

# Description of Changes

Finish the work started in #3227, which switched a hardcoded `127.0.0.1`
to a dynamic server address.

# API and ABI breaking changes

None.

# Expected complexity level and risk

1

# Testing

- [x] Smoketests still pass
- [ ] `Internal Tests` pass under this PR (they don't on master)

---------

Co-authored-by: Zeke Foppa <[email protected]>

Author: bfops

smoketests/__init__.py

@@ -301,18 +301,22 @@ def get_server_address(self):
             server_config = next((c for c in config['server_configs'] if c['nickname'] == server_name), None)
             if server_config is None:
                 raise Exception(f"Unable to find server in config with nickname {server_name}")
-            host = server_config['host']
+            address = server_config['host']
+            host = address
+            port = None
+            if ":" in host:
+                host, port = host.split(":", 1)
             protocol = server_config['protocol']
 
-            return dict(host=host, protocol=protocol, token=token)
+            return dict(address=address, host=host, port=port, protocol=protocol, token=token)
 
     # Make an HTTP call with `method` to `path`.
     #
     # If the response is 200, return the body.
     # Otherwise, throw an `Exception` constructed with two arguments, the response object and the body.
     def api_call(self, method, path, body=None, headers={}):
         server = self.get_server_address()
-        host = server["host"]
+        host = server["address"]
         protocol = server["protocol"]
         token = server["token"]
         conn = None

smoketests/tests/pg_wire.py

@@ -4,30 +4,6 @@
 import tomllib
 import psycopg2
 
-
-def psql(identity: str, sql: str) -> str:
-    """Call `psql` and execute the given SQL statement."""
-    result = subprocess.run(
-        ["psql", "-h", "127.0.0.1", "-p", "5432", "-U", "postgres", "-d", "quickstart", "--quiet", "-c", sql],
-        encoding="utf8",
-        env={**os.environ, "PGPASSWORD": identity},
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
-        text=True,
-    )
-
-    if result.stderr:
-        raise Exception(result.stderr.strip())
-    return result.stdout.strip()
-
-
-def connect_db(identity: str):
-    """Connect to the database using `psycopg2`."""
-    conn = psycopg2.connect(host="127.0.0.1", port=5432, user="postgres", password=identity, dbname="quickstart")
-    conn.set_session(autocommit=True)  # Disable automic transaction
-    return conn
-
-
 class SqlFormat(Smoketest):
     AUTOPUBLISH = False
     MODULE_CODE = """
@@ -168,9 +144,32 @@ class SqlFormat(Smoketest):
 }
 """
 
+    def psql(self, identity: str, sql: str) -> str:
+        """Call `psql` and execute the given SQL statement."""
+        server = self.get_server_address()
+        result = subprocess.run(
+            ["psql", "-h", server["host"], "-p", "5432", "-U", "postgres", "-d", "quickstart", "--quiet", "-c", sql],
+            encoding="utf8",
+            env={**os.environ, "PGPASSWORD": identity},
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+
+        if result.stderr:
+            raise Exception(result.stderr.strip())
+        return result.stdout.strip()
+
+    def connect_db(self, identity: str):
+        """Connect to the database using `psycopg2`."""
+        server = self.get_server_address()
+        conn = psycopg2.connect(host=server["host"], port=5432, user="postgres", password=identity, dbname="quickstart")
+        conn.set_session(autocommit=True)  # Disable automic transaction
+        return conn
+
     def assertSql(self, token: str, sql: str, expected):
         self.maxDiff = None
-        sql_out = psql(token, sql)
+        sql_out = self.psql(token, sql)
         sql_out = "\n".join([line.rstrip() for line in sql_out.splitlines()])
         expected = "\n".join([line.rstrip() for line in expected.splitlines()])
         print(sql_out)
@@ -242,7 +241,7 @@ def test_sql_conn(self):
         self.publish_module("quickstart", clear=True)
         self.call("test")
 
-        conn = connect_db(token)
+        conn = self.connect_db(token)
         # Check prepared statements (faked by `psycopg2`)
         with conn.cursor() as cur:
             cur.execute("select * from t_uints where u8 = %s and u16 = %s", (105, 1050))
@@ -262,20 +261,20 @@ def test_failures(self):
         self.publish_module("quickstart", clear=True)
 
         # Empty query
-        sql_out = psql(token, "")
+        sql_out = self.psql(token, "")
         self.assertEqual(sql_out, "")
 
         # Connection fails with invalid token
         with self.assertRaises(Exception) as cm:
-            psql("invalid_token", "SELECT * FROM t_uints")
+            self.psql("invalid_token", "SELECT * FROM t_uints")
         self.assertIn("Invalid token", str(cm.exception))
 
         # Returns error for unsupported `sql` statements
         with self.assertRaises(Exception) as cm:
-            psql(token, "SELECT CASE a WHEN 1 THEN 'one' ELSE 'other' END FROM t_uints")
+            self.psql(token, "SELECT CASE a WHEN 1 THEN 'one' ELSE 'other' END FROM t_uints")
         self.assertIn("Unsupported", str(cm.exception))
 
         # And prepared statements
         with self.assertRaises(Exception) as cm:
-            psql(token, "SELECT * FROM t_uints where u8 = $1")
+            self.psql(token, "SELECT * FROM t_uints where u8 = $1")
         self.assertIn("Unsupported", str(cm.exception))

smoketests/tests/quickstart.py

@@ -125,7 +125,7 @@ def _test_quickstart(self):
             main = main.replace(src, dst)
         main += "\n" + self.extra_code
         server = self.get_server_address()
-        host = server["host"]
+        host = server["address"]
         protocol = server["protocol"]
         main = main.replace("http://localhost:3000", f"{protocol}://{host}")
         _write_file(client_path / self.client_file, main)