From 9e61284d104d245f93b2b6feb7538e7246a1c677 Mon Sep 17 00:00:00 2001
From: Alexey Masterov <alexeymasterov@neon.tech>
Date: Thu, 5 Sep 2024 11:55:23 +0200
Subject: [PATCH] fix mypy warnings

---
 .../cloud_regress/test_cloud_regress.py       | 108 +++++++++---------
 1 file changed, 57 insertions(+), 51 deletions(-)

diff --git a/test_runner/cloud_regress/test_cloud_regress.py b/test_runner/cloud_regress/test_cloud_regress.py
index dcd1898afa..708f9203a7 100644
--- a/test_runner/cloud_regress/test_cloud_regress.py
+++ b/test_runner/cloud_regress/test_cloud_regress.py
@@ -18,56 +18,62 @@ def test_cloud_regress(remote_pg: RemotePostgres):
     """
     Run the regression tests
     """
-    pg_version = re.search(r"\-pg(\d+)\]", os.environ.get("PYTEST_CURRENT_TEST")).group(1)
-    conn = psycopg2.connect(remote_pg.connstr())
-    cur = conn.cursor()
-    cur.execute("SELECT COUNT(*) FROM pg_extension WHERE extname = 'regress_so'")
-    num_ext = cur.fetchone()[0]
-    assert int(num_ext) < 2
-    if num_ext == 1:
-        log.info("The extension is found")
-    else:
-        log.info("Creating the extension")
-        cur.execute("CREATE EXTENSION regress_so")
-        conn.commit()
+    cur_test = os.environ.get("PYTEST_CURRENT_TEST")
+    assert cur_test
+    pg_version_match = re.match(r"\-pg(\d+)\]", cur_test)
+    assert pg_version_match
+    pg_version = int(pg_version_match.group(1))
+    with psycopg2.connect(remote_pg.connstr()) as conn:
+        with conn.cursor() as cur:
+            cur = conn.cursor()
+            cur.execute("SELECT COUNT(*) FROM pg_extension WHERE extname = 'regress_so'")
+            for r in cur:
+                num_ext = r[0]
+            assert int(num_ext) < 2
+            if num_ext == 1:
+                log.info("The extension is found")
+            else:
+                log.info("Creating the extension")
+                cur.execute("CREATE EXTENSION regress_so")
+                conn.commit()
 
-    log.info("Creating a C function to check availability of regress.so")
-    cur.execute(
-        "CREATE FUNCTION get_columns_length(oid[]) "
-        "RETURNS int AS 'regress.so' LANGUAGE C STRICT STABLE PARALLEL SAFE;"
-    )
-    conn.rollback()
-    neondir = os.path.abspath(os.path.join(os.path.dirname(os.path.relpath(__file__)), "../../"))
-    runpath = f"{neondir}/vendor/postgres-v{pg_version}/src/test/regress"
-    artpath = "/tmp/neon"
-    prefix = f"{artpath}/pg_install/v{pg_version}"
-    binpath = f"{prefix}/bin"
-    regress_bin = f"{prefix}/lib/postgresql/pgxs/src/test/regress/pg_regress"
+            log.info("Creating a C function to check availability of regress.so")
+            cur.execute(
+                "CREATE FUNCTION get_columns_length(oid[]) "
+                "RETURNS int AS 'regress.so' LANGUAGE C STRICT STABLE PARALLEL SAFE;"
+            )
+            conn.rollback()
+            neondir = os.path.abspath(os.path.join(os.path.dirname(os.path.relpath(__file__)), "../../"))
+            runpath = f"{neondir}/vendor/postgres-v{pg_version}/src/test/regress"
+            artpath = "/tmp/neon"
+            prefix = f"{artpath}/pg_install/v{pg_version}"
+            binpath = f"{prefix}/bin"
+            regress_bin = f"{prefix}/lib/postgresql/pgxs/src/test/regress/pg_regress"
 
-    env_vars = {
-        "PGHOST": remote_pg.default_options["host"],
-        "PGPORT": str(
-            remote_pg.default_options["port"] if "port" in remote_pg.default_options else 5432
-        ),
-        "PGUSER": remote_pg.default_options["user"],
-        "PGPASSWORD": remote_pg.default_options["password"],
-        "PGDATABASE": remote_pg.default_options["dbname"],
-    }
-    regress_cmd = [
-        regress_bin,
-        "--inputdir=.",
-        f"--bindir={binpath}",
-        "--dlpath=/usr/local/lib",
-        "--max-concurrent-tests=20",
-        "--schedule=./parallel_schedule",
-        "--max-connections=5",
-    ]
-    try:
-        remote_pg.pg_bin.run(regress_cmd, env=env_vars, cwd=runpath)
-    except subprocess.CalledProcessError as e:
-        log.error("Error(s) occurred while running the regression tests")
-        with open(f"{runpath}/regression.out", "r") as f:
-            print(f.read())
-        with open(f"{runpath}/regression.diffs", "r") as f:
-            print(f.read())
-        raise e
+            env_vars = {
+                "PGHOST": remote_pg.default_options["host"],
+                "PGPORT": str(
+                    remote_pg.default_options["port"] if "port" in remote_pg.default_options else 5432
+                ),
+                "PGUSER": remote_pg.default_options["user"],
+                "PGPASSWORD": remote_pg.default_options["password"],
+                "PGDATABASE": remote_pg.default_options["dbname"],
+            }
+            regress_cmd = [
+                regress_bin,
+                "--inputdir=.",
+                f"--bindir={binpath}",
+                "--dlpath=/usr/local/lib",
+                "--max-concurrent-tests=20",
+                "--schedule=./parallel_schedule",
+                "--max-connections=5",
+            ]
+            try:
+                remote_pg.pg_bin.run(regress_cmd, env=env_vars, cwd=runpath)
+            except subprocess.CalledProcessError as e:
+                log.error("Error(s) occurred while running the regression tests")
+                with open(f"{runpath}/regression.out", "r") as f:
+                    print(f.read())
+                with open(f"{runpath}/regression.diffs", "r") as f:
+                    print(f.read())
+                raise e