From 986ee663589690c78ed52dcf45334d6eb5621bad Mon Sep 17 00:00:00 2001
From: Anastasia Lubennikova <anastasia@neon.tech>
Date: Mon, 24 Jul 2023 16:46:21 +0300
Subject: [PATCH] Improve test_download_extensions: - fix cleanup; - add test
 LOAD library

---
 .../regress/test_download_extensions.py       | 96 ++++++++++++++++++-
 1 file changed, 94 insertions(+), 2 deletions(-)

diff --git a/test_runner/regress/test_download_extensions.py b/test_runner/regress/test_download_extensions.py
index 4a5729df04..d3c1415ec3 100644
--- a/test_runner/regress/test_download_extensions.py
+++ b/test_runner/regress/test_download_extensions.py
@@ -75,22 +75,24 @@ def test_remote_extensions(
             assert "embedding" in [x[0] for x in cur.fetchall()]
 
             # check that we can download private extension
+            # this is expected to fail because we 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"
 
-            # TODO: try to load libraries as well
-
     # Cleaning up downloaded files is important for local tests
     # or else one test could reuse the files from another test or another test run
     cleanup_files = [
         "lib/postgresql/anon.so",
         "lib/postgresql/embedding.so",
         "share/postgresql/extension/anon.control",
+        "share/postgresql/extension/anon.control.zst",
         "share/postgresql/extension/embedding--0.1.0.sql",
+        "share/postgresql/extension/embedding--0.1.0.sql.zst",
         "share/postgresql/extension/embedding.control",
+        "share/postgresql/extension/embedding.control.zst",
     ]
     cleanup_files = [f"pg_install/v{pg_version}/" + x for x in cleanup_files]
     cleanup_folders = [
@@ -98,15 +100,105 @@ def test_remote_extensions(
         f"pg_install/v{pg_version}/share/postgresql/extension/anon",
         f"pg_install/v{pg_version}/extensions",
     ]
+
     for file in cleanup_files:
         try:
             os.remove(file)
             log.info(f"removed file {file}")
         except Exception as err:
             log.info(f"error removing file {file}: {err}")
+
     for folder in cleanup_folders:
         try:
             shutil.rmtree(folder)
             log.info(f"removed folder {folder}")
+        except Exception as err:
+            log.info(f"error removing folder {folder}: {err}")
+
+
+# Test downloading remote library.
+@pytest.mark.parametrize("remote_storage_kind", available_s3_storages())
+def test_remote_library(
+    neon_env_builder: NeonEnvBuilder,
+    remote_storage_kind: RemoteStorageKind,
+    pg_version: PgVersion,
+):
+    neon_env_builder.enable_remote_storage(
+        remote_storage_kind=remote_storage_kind,
+        test_name="test_remote_library",
+        enable_remote_extensions=True,
+    )
+    neon_env_builder.num_safekeepers = 3
+    env = neon_env_builder.init_start()
+    tenant_id, _ = env.neon_cli.create_tenant()
+    env.neon_cli.create_timeline("test_remote_library", tenant_id=tenant_id)
+
+    assert env.ext_remote_storage is not None  # satisfy mypy
+    assert env.remote_storage_client is not None  # satisfy mypy
+
+    # 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:]
+
+                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",
+        tenant_id=tenant_id,
+        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
+
+    # Cleaning up downloaded files is important for local tests
+    # or else one test could reuse the files from another test or another test run
+    cleanup_files = [
+        "lib/postgresql/anon.so",
+        "share/postgresql/extension/anon.control",
+        "share/postgresql/extension/anon.control.zst",
+    ]
+    cleanup_files = [f"pg_install/v{pg_version}/" + x for x in cleanup_files]
+    cleanup_folders = [
+        "extensions",
+        f"pg_install/v{pg_version}/share/postgresql/extension/anon",
+        f"pg_install/v{pg_version}/extensions",
+    ]
+
+    for file in cleanup_files:
+        try:
+            os.remove(file)
+            log.info(f"removed file {file}")
         except Exception as err:
             log.info(f"error removing file {file}: {err}")
+
+    for folder in cleanup_folders:
+        try:
+            shutil.rmtree(folder)
+            log.info(f"removed folder {folder}")
+        except Exception as err:
+            log.info(f"error removing folder {folder}: {err}")
+
+
+# TODO test library which name is different from extension name