From 11531b5e8396f6c4f28da3ad3f63ce42db57f5fb Mon Sep 17 00:00:00 2001 From: Alek Westover Date: Thu, 27 Jul 2023 14:51:39 -0400 Subject: [PATCH] Making the test code a little prettier --- .../regress/test_download_extensions.py | 174 ++++++++---------- 1 file changed, 78 insertions(+), 96 deletions(-) diff --git a/test_runner/regress/test_download_extensions.py b/test_runner/regress/test_download_extensions.py index 001ef00919..1b3ae7f992 100644 --- a/test_runner/regress/test_download_extensions.py +++ b/test_runner/regress/test_download_extensions.py @@ -18,7 +18,10 @@ def add_pgdir_prefix(pgversion, files): # Cleaning up downloaded files is important for local tests # or else one test could reuse the files from another test or another test run -def cleanup(cleanup_files, cleanup_folders): +def cleanup(cleanup_files, cleanup_folders, pg_version): + cleanup_files = add_pgdir_prefix(pg_version, cleanup_files) + cleanup_folders = add_pgdir_prefix(pg_version, cleanup_folders) + for file in cleanup_files: try: os.remove(file) @@ -34,6 +37,32 @@ def cleanup(cleanup_files, cleanup_folders): log.info(f"error removing folder {folder}: {err}") +cleanup_files = [ + "lib/postgresql/anon.so", + "share/postgresql/extension/anon.control", +] +cleanup_folders = ["share/postgresql/extension/anon", "download_extensions"] + + +def upload_files(env): + log.info("Uploading test files to mock bucket") + os.chdir("test_runner/regress/data/extension_test") + for path in os.walk("."): + prefix, _, files = path + for file in files: + # the [2:] is to remove the leading "./" + full_path = os.path.join(prefix, file)[2:] + + with open(full_path, "rb") as f: + log.info(f"UPLOAD {full_path} to ext/{full_path}") + env.remote_storage_client.upload_fileobj( + f, + env.ext_remote_storage.bucket_name, + f"ext/{full_path}", + ) + os.chdir("../../../..") + + # Test downloading remote extension. @pytest.mark.parametrize("remote_storage_kind", available_s3_storages()) def test_remote_extensions( @@ -57,22 +86,8 @@ def test_remote_extensions( # For MOCK_S3 we upload test files. # For REAL_S3 we use the files already in the bucket if remote_storage_kind == RemoteStorageKind.MOCK_S3: - log.info("Uploading test files to mock bucket") - os.chdir("test_runner/regress/data/extension_test") - for path in os.walk("."): - prefix, _, files = path - for file in files: - # the [2:] is to remove the leading "./" - full_path = os.path.join(prefix, file)[2:] + upload_files(env) - with open(full_path, "rb") as f: - log.info(f"UPLOAD {full_path} to ext/{full_path}") - env.remote_storage_client.upload_fileobj( - f, - env.ext_remote_storage.bucket_name, - f"ext/{full_path}", - ) - os.chdir("../../../..") # Start a compute node and check that it can download the extensions # and use them to CREATE EXTENSION and LOAD endpoint = env.endpoints.create_start( @@ -81,43 +96,35 @@ def test_remote_extensions( remote_ext_config=env.ext_remote_storage.to_string(), # config_lines=["log_min_messages=debug3"], ) - with closing(endpoint.connect()) as conn: - with conn.cursor() as cur: - # Check that appropriate control files were downloaded - cur.execute("SELECT * FROM pg_available_extensions") - all_extensions = [x[0] for x in cur.fetchall()] - log.info(all_extensions) - assert "anon" in all_extensions - assert "kq_imcx" in all_extensions + try: + with closing(endpoint.connect()) as conn: + with conn.cursor() as cur: + # Check that appropriate control files were downloaded + cur.execute("SELECT * FROM pg_available_extensions") + all_extensions = [x[0] for x in cur.fetchall()] + log.info(all_extensions) + assert "anon" in all_extensions + assert "kq_imcx" in all_extensions - # postgis is on real s3 but not mock s3. - # it's kind of a big file, would rather not upload to github - if remote_storage_kind == RemoteStorageKind.REAL_S3: - assert "postgis" in all_extensions - # this is expected to break on my computer because I lack the necesary dependencies + # postgis is on real s3 but not mock s3. + # it's kind of a big file, would rather not upload to github + if remote_storage_kind == RemoteStorageKind.REAL_S3: + assert "postgis" in all_extensions + # this is expected to break on my computer because I lack the necesary dependencies + try: + cur.execute("CREATE EXTENSION postgis") + except Exception as err: + log.info(f"(expected) error creating postgis extension: {err}") + + # this is expected to fail on my computer because I don't have the pgcrypto extension try: - cur.execute("CREATE EXTENSION postgis") + cur.execute("CREATE EXTENSION anon") except Exception as err: - log.info(f"(expected) error creating postgis extension: {err}") - - # this is expected to fail on my computer because I don't have the pgcrypto extension - try: - cur.execute("CREATE EXTENSION anon") - except Exception as err: - log.info("error creating anon extension") - assert "pgcrypto" in str(err), "unexpected error creating anon extension" - - cleanup_files = add_pgdir_prefix( - pg_version, - [ - "lib/postgresql/anon.so", - "share/postgresql/extension/anon.control", - ], - ) - cleanup_folders = add_pgdir_prefix( - pg_version, ["share/postgresql/extension/anon", "download_extensions"] - ) - cleanup(cleanup_files, cleanup_folders) + log.info("error creating anon extension") + assert "pgcrypto" in str(err), "unexpected error creating anon extension" + finally: + pass + # cleanup(cleanup_files, cleanup_folders, pg_version) # Test downloading remote library. @@ -143,23 +150,8 @@ def test_remote_library( # For MOCK_S3 we upload test files. # For REAL_S3 we use the files already in the bucket if remote_storage_kind == RemoteStorageKind.MOCK_S3: - log.info("Uploading test files to mock bucket") - os.chdir("test_runner/regress/data/extension_test") - for path in os.walk("."): - prefix, _, files = path - for file in files: - # the [2:] is to remove the leading "./" - full_path = os.path.join(prefix, file)[2:] + upload_files(env) - with open(full_path, "rb") as f: - log.info(f"UPLOAD {full_path} to ext/{full_path}") - env.remote_storage_client.upload_fileobj( - f, - env.ext_remote_storage.bucket_name, - f"ext/{full_path}", - ) - os.chdir("../../../..") - # Start a compute node and check that it can download the extensions # and use them to run LOAD library endpoint = env.endpoints.create_start( "test_remote_library", @@ -167,36 +159,26 @@ def test_remote_library( remote_ext_config=env.ext_remote_storage.to_string(), # config_lines=["log_min_messages=debug3"], ) - with closing(endpoint.connect()) as conn: - with conn.cursor() as cur: - # try to load library - try: - cur.execute("LOAD 'anon'") - except Exception as err: - log.info(f"error loading anon library: {err}") - raise AssertionError("unexpected error loading anon library") from err - - # test library which name is different from extension name - # this fails on my computer because I' missing a dependency - # however, it does successfully download the postgis archive - if remote_storage_kind == RemoteStorageKind.REAL_S3: + try: + with closing(endpoint.connect()) as conn: + with conn.cursor() as cur: + # try to load library try: - cur.execute("LOAD 'postgis_topology-3'") + cur.execute("LOAD 'anon'") except Exception as err: - log.info("error loading postgis_topology-3") - assert ( - "libproj.so.19: cannot open shared object file: No such file or directory" - in str(err) - ), "unexpected error loading postgis_topology-3" + log.info(f"error loading anon library: {err}") + raise AssertionError("unexpected error loading anon library") from err - cleanup_files = add_pgdir_prefix( - pg_version, - [ - "lib/postgresql/anon.so", - "share/postgresql/extension/anon.control", - ], - ) - cleanup_folders = add_pgdir_prefix( - pg_version, ["share/postgresql/extension/anon", "download_extensions"] - ) - cleanup(cleanup_files, cleanup_folders) + # test library which name is different from extension name + # this fails on my computer because I' missing a dependency + # however, it does successfully download the postgis archive + if remote_storage_kind == RemoteStorageKind.REAL_S3: + try: + cur.execute("LOAD 'postgis_topology-3'") + except Exception as err: + log.info("error loading postgis_topology-3") + assert "cannot open shared object file: No such file or directory" in str( + err + ), "unexpected error loading postgis_topology-3" + finally: + cleanup(cleanup_files, cleanup_folders, pg_version)