This is an automated email from the ASF dual-hosted git repository.
yuqi1129 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gravitino.git
The following commit(s) were added to refs/heads/main by this push:
new d4cad80ee9 [#13323] test(client): Add Ray Data Iceberg REST contract
tests (#13325)
d4cad80ee9 is described below
commit d4cad80ee938e59fc872dc77e002e275d92d72c1
Author: Thunderkeg <[email protected]>
AuthorDate: Tue Oct 27 09:35:04 2026 +0800
[#13323] test(client): Add Ray Data Iceberg REST contract tests (#13325)
### What changes were proposed in this pull request?
This pull request adds a Gravitino-owned Ray Data Iceberg REST
integration test to the Python client test suite and runs it in a
dedicated Python integration CI job.
The change:
- adds pinned Ray 2.55.1, PyIceberg 0.12.0, and pandas 2.3.3 test
dependencies;
- adds a dedicated `rayIcebergIT` Gradle task with an isolated Python
environment;
- generates the ignored `gravitino/version.ini` file before running the
client test;
- configures the Gravitino Iceberg REST auxiliary service with the
dynamic catalog provider;
- uses Gravitino's memory Iceberg catalog and PyIceberg to create a
namespace and pre-existing table fixture;
- verifies empty-table reads, Ray Data writes, subsequent reads, field
order, row counts, and filtering;
- adds setup-failure cleanup for Gravitino configuration, metadata, and
temporary data;
- runs `rayIcebergIT` in a dedicated Python 3.12 CI job when the
existing Python integration source-change filter matches;
- makes the dedicated task fail instead of skip when Ray or PyIceberg is
unavailable.
The implementation does not modify Java Iceberg REST code, Ray source
code, or Ray `Catalog.resolve()` behavior.
### Why are the changes needed?
Gravitino documents Ray Data access through the Iceberg REST endpoint
and `catalog_kwargs`, but the Python client integration suite did not
have a real end-to-end regression test for this contract.
The test protects the existing-table read/write path and validates the
Gravitino auxiliary service lifecycle using the existing Lance Ray IT as
the implementation baseline; wiring the task into CI ensures the
contract is exercised automatically.
Fix: #13323
### Does this PR introduce _any_ user-facing change?
No end-user API or runtime behavior changes are introduced;
contributor-facing Python integration-test tooling is added.
### How was this patch tested?
Passed:
- `git diff --check`
- `ruff check --no-cache
clients/client-python/tests/integration/test_ray_iceberg.py`
- `ruff format --check --no-cache
clients/client-python/tests/integration/test_ray_iceberg.py`
- `actionlint .github/workflows/python-integration-test.yml
.github/workflows/ray-iceberg-it-action.yml`
- Gradle task graph for `:clients:client-python:rayIcebergIT --dry-run`
- `./gradlew --offline -PpythonVersion=3.12 -PskipDockerTests=true
-PskipWeb=true :clients:client-python:rayIcebergIT --no-daemon
--console=plain` — passed; 1 test ran, 0 skipped.
---------
Signed-off-by: jiangxt2 <[email protected]>
---
.github/workflows/python-integration-test.yml | 11 +
.github/workflows/ray-iceberg-it-action.yml | 71 ++++++
clients/client-python/build.gradle.kts | 61 +++++
clients/client-python/requirements-ray-iceberg.txt | 21 ++
.../tests/integration/integration_test_env.py | 16 +-
.../tests/integration/test_ray_iceberg.py | 265 +++++++++++++++++++++
.../tests/unittests/test_integration_test_env.py | 146 ++++++++++++
7 files changed, 587 insertions(+), 4 deletions(-)
diff --git a/.github/workflows/python-integration-test.yml
b/.github/workflows/python-integration-test.yml
index 653ee4e8fb..4670c221ce 100644
--- a/.github/workflows/python-integration-test.yml
+++ b/.github/workflows/python-integration-test.yml
@@ -92,3 +92,14 @@ jobs:
sync: ${{ inputs.sync }}
architecture: ${{ matrix.architecture }}
java-version: ${{ matrix.java-version }}
+
+ RayIcebergIT:
+ needs: changes
+ if: needs.changes.outputs.source_changes == 'true'
+ uses: ./.github/workflows/ray-iceberg-it-action.yml
+ with:
+ base-ref: ${{ inputs.base-ref }}
+ base-sha: ${{ inputs.base-sha }}
+ head-repository: ${{ inputs.head-repository }}
+ head-sha: ${{ inputs.head-sha }}
+ sync: ${{ inputs.sync }}
diff --git a/.github/workflows/ray-iceberg-it-action.yml
b/.github/workflows/ray-iceberg-it-action.yml
new file mode 100644
index 0000000000..06278a576b
--- /dev/null
+++ b/.github/workflows/ray-iceberg-it-action.yml
@@ -0,0 +1,71 @@
+name: Ray Iceberg REST Contract Test Action
+
+on:
+ workflow_call:
+ inputs:
+ base-ref:
+ type: string
+ required: false
+ default: ''
+ base-sha:
+ type: string
+ required: false
+ default: ''
+ head-repository:
+ type: string
+ required: false
+ default: ''
+ head-sha:
+ type: string
+ required: false
+ default: ''
+ sync:
+ type: string
+ required: false
+ default: ''
+
+jobs:
+ RayIcebergIT:
+ name: Ray Iceberg REST contract (Python 3.12)
+ runs-on: ubuntu-latest
+ timeout-minutes: 120
+ steps:
+ - uses: actions/checkout@v4
+ with:
+ persist-credentials: false
+ - uses: ./.github/actions/checkout-and-sync
+ with:
+ base-ref: ${{ inputs.base-ref }}
+ base-sha: ${{ inputs.base-sha }}
+ head-repository: ${{ inputs.head-repository }}
+ head-sha: ${{ inputs.head-sha }}
+ sync: ${{ inputs.sync }}
+
+ - uses: actions/setup-java@v4
+ with:
+ java-version: 17
+ distribution: "temurin"
+ cache: "gradle"
+
+ - name: Free up disk space
+ run: dev/ci/util_free_space.sh
+
+ - name: Package Gravitino
+ run: ./gradlew compileDistribution -PskipWeb=true -x test
+
+ - name: Run Ray Iceberg REST contract test
+ id: rayIcebergIT
+ run: |
+ ./gradlew -PpythonVersion=3.12 -PskipDockerTests=true -PskipWeb=true
\
+ :clients:client-python:rayIcebergIT
+
+ - name: Upload Ray Iceberg IT reports and logs
+ uses: actions/upload-artifact@v7
+ if: ${{ failure() && steps.rayIcebergIT.outcome == 'failure' }}
+ with:
+ name: ray-iceberg-it-reports
+ path: |
+ build/reports
+ integration-test/build/integration-test.log
+ distribution/package/logs/gravitino-server.out
+ distribution/package/logs/gravitino-server.log
diff --git a/clients/client-python/build.gradle.kts
b/clients/client-python/build.gradle.kts
index 3311608c80..41de64f8dd 100644
--- a/clients/client-python/build.gradle.kts
+++ b/clients/client-python/build.gradle.kts
@@ -226,6 +226,67 @@ tasks {
environment = envMap
}
+ val rayIcebergPythonVersion =
project.rootProject.extra["pythonVersion"].toString()
+ val rayIcebergOsDir = when {
+ org.gradle.internal.os.OperatingSystem.current().isMacOsX -> "MacOSX"
+ org.gradle.internal.os.OperatingSystem.current().isLinux -> "Linux"
+ else -> throw GradleException("Ray Iceberg IT only supports macOS and
Linux")
+ }
+ val rayIcebergMiniforgeDir =
+
project.file("${project.rootDir}/.gradle/python/$rayIcebergOsDir/Miniforge3")
+ val rayIcebergCondaExecutable = rayIcebergMiniforgeDir.resolve("bin/conda")
+ val rayIcebergEnvDir =
+ project.layout.buildDirectory.dir("ray-iceberg-it-env").get().asFile
+ val rayIcebergPythonExecutable = rayIcebergEnvDir.resolve("bin/python")
+
+ val rayIcebergDependencies by registering {
+ group = "verification"
+ description = "Install the dependencies for the Ray Iceberg IT."
+ dependsOn("miniforgeSetup")
+ doLast {
+ if (!rayIcebergPythonExecutable.exists()) {
+ project.exec {
+ executable = rayIcebergCondaExecutable.absolutePath
+ args = listOf(
+ "create", "--prefix", rayIcebergEnvDir.absolutePath,
+ "python=$rayIcebergPythonVersion", "--yes"
+ )
+ }
+ }
+ project.exec {
+ executable = rayIcebergPythonExecutable.absolutePath
+ workingDir = projectDir
+ args = listOf(
+ "-m", "pip", "install", "-e", ".", "-r",
"requirements-ray-iceberg.txt"
+ )
+ }
+ project.exec {
+ executable = rayIcebergPythonExecutable.absolutePath
+ workingDir = projectDir
+ args = listOf("scripts/generate_version.py")
+ }
+ }
+ }
+
+ register("rayIcebergIT") {
+ group = "verification"
+ description = "Run the Gravitino Iceberg REST contract test with Ray Data."
+ dependsOn(rayIcebergDependencies, startGravitinoServer)
+ finalizedBy(stopGravitinoServer)
+ doLast {
+ project.exec {
+ executable = rayIcebergPythonExecutable.absolutePath
+ workingDir = projectDir.resolve("./tests/integration")
+ environment("PROJECT_VERSION", project.version)
+ environment("GRAVITINO_HOME", project.rootDir.path +
"/distribution/package")
+ environment("START_EXTERNAL_GRAVITINO", "true")
+ environment("RAY_ICEBERG_IT_REQUIRED", "true")
+ environment("PYTHONPATH",
"${project.rootDir.path}/clients/client-python")
+ args = listOf("-m", "unittest", "-v", "test_ray_iceberg")
+ }
+ }
+ }
+
val unitCoverageReport by registering(VenvTask::class){
venvExec = "coverage"
args = listOf("html")
diff --git a/clients/client-python/requirements-ray-iceberg.txt
b/clients/client-python/requirements-ray-iceberg.txt
new file mode 100644
index 0000000000..5d61395414
--- /dev/null
+++ b/clients/client-python/requirements-ray-iceberg.txt
@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+ray[data]==2.55.1
+pyiceberg==0.12.0
+pandas==2.3.3
+pyarrow==23.0.1
diff --git a/clients/client-python/tests/integration/integration_test_env.py
b/clients/client-python/tests/integration/integration_test_env.py
index 06fa23588d..1fe656f1a1 100644
--- a/clients/client-python/tests/integration/integration_test_env.py
+++ b/clients/client-python/tests/integration/integration_test_env.py
@@ -259,17 +259,24 @@ class IntegrationTestEnv(unittest.TestCase):
"project root directory."
)
+ env_vars = os.environ.copy()
+ env_vars["HADOOP_USER_NAME"] = "anonymous"
+ # Stop the server before removing data so it cannot recreate files
during cleanup.
+ subprocess.run(
+ [gravitino_startup_script, "stop"],
+ env=env_vars,
+ check=True,
+ )
+
# remove data dir under gravitino_home
data_dir = os.path.join(gravitino_home, "data")
if os.path.exists(data_dir):
logger.info("Remove Gravitino data directory: %s", data_dir)
shutil.rmtree(data_dir)
- # Restart Gravitino Server
- env_vars = os.environ.copy()
- env_vars["HADOOP_USER_NAME"] = "anonymous"
+ # Start Gravitino Server
result = subprocess.run(
- [gravitino_startup_script, "restart"],
+ [gravitino_startup_script, "start"],
env=env_vars,
capture_output=True,
text=True,
@@ -279,6 +286,7 @@ class IntegrationTestEnv(unittest.TestCase):
logger.info("stdout: %s", result.stdout)
if result.stderr:
logger.info("stderr: %s", result.stderr)
+ result.check_returncode()
success = check_gravitino_server_status()
if success:
diff --git a/clients/client-python/tests/integration/test_ray_iceberg.py
b/clients/client-python/tests/integration/test_ray_iceberg.py
new file mode 100644
index 0000000000..3e79373edf
--- /dev/null
+++ b/clients/client-python/tests/integration/test_ray_iceberg.py
@@ -0,0 +1,265 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import logging
+import os
+import shutil
+import tempfile
+import time
+import unittest
+from typing import Optional
+from uuid import uuid4
+
+import requests
+
+from gravitino import Catalog, GravitinoAdminClient, GravitinoClient
+from tests.integration.integration_test_env import IntegrationTestEnv
+
+logger = logging.getLogger(__name__)
+
+ICEBERG_REST_BASE_URL = "http://localhost:9001/iceberg/";
+MAIN_CONF_FILE = "conf/gravitino.conf"
+ICEBERG_REST_CONFIG_PROVIDER_KEY =
"gravitino.iceberg-rest.catalog-config-provider"
+ICEBERG_REST_METALAKE_KEY = "gravitino.iceberg-rest.gravitino-metalake"
+ICEBERG_REST_DEFAULT_CATALOG_KEY =
"gravitino.iceberg-rest.default-catalog-name"
+
+
+def _missing_ray_iceberg_dependencies() -> Optional[str]:
+ missing = []
+ for module in ("ray", "pyiceberg"):
+ try:
+ __import__(module)
+ except ImportError:
+ missing.append(module)
+ return ", ".join(missing) if missing else None
+
+
+_MISSING_RAY_ICEBERG_DEPS = _missing_ray_iceberg_dependencies()
+_RAY_ICEBERG_IT_REQUIRED = os.environ.get("RAY_ICEBERG_IT_REQUIRED") == "true"
+
+
[email protected](
+ _MISSING_RAY_ICEBERG_DEPS is not None and not _RAY_ICEBERG_IT_REQUIRED,
+ f"Ray Iceberg test deps not installed: {_MISSING_RAY_ICEBERG_DEPS}. "
+ "Install with: pip install -r
clients/client-python/requirements-ray-iceberg.txt",
+)
+class TestRayIcebergIntegration(IntegrationTestEnv):
+ """Exercise Ray Data's Iceberg REST read/write contract with Gravitino."""
+
+ METALAKE_NAME = f"ray_iceberg_it_{uuid4().hex[:8]}"
+ CATALOG_NAME = f"ray_iceberg_{uuid4().hex[:8]}"
+ SCHEMA_NAME = "ray_schema"
+ TABLE_NAME = "ray_table"
+ TABLE_IDENTIFIER = f"{SCHEMA_NAME}.{TABLE_NAME}"
+
+ gravitino_admin_client: Optional[GravitinoAdminClient] = None
+ gravitino_client: Optional[GravitinoClient] = None
+ iceberg_catalog = None
+ temp_dir: Optional[str] = None
+ main_conf_path: Optional[str] = None
+ appended_iceberg_rest_conf = False
+ setup_completed = False
+
+ @classmethod
+ def setUpClass(cls):
+ cls.addClassCleanup(cls._cleanup_after_setup_failure)
+ super().setUpClass()
+
+ cls._get_gravitino_home()
+ cls.main_conf_path = os.path.join(cls.gravitino_home, MAIN_CONF_FILE)
+ cls._append_conf(
+ {
+ ICEBERG_REST_CONFIG_PROVIDER_KEY: "dynamic-config-provider",
+ ICEBERG_REST_METALAKE_KEY: cls.METALAKE_NAME,
+ ICEBERG_REST_DEFAULT_CATALOG_KEY: cls.CATALOG_NAME,
+ },
+ cls.main_conf_path,
+ )
+ cls.appended_iceberg_rest_conf = True
+ cls.restart_server()
+
+ cls.gravitino_admin_client =
GravitinoAdminClient("http://localhost:8090";)
+ cls.gravitino_admin_client.create_metalake(
+ cls.METALAKE_NAME,
+ comment="Ray Iceberg IT metalake",
+ properties={},
+ )
+ cls.gravitino_client = GravitinoClient(
+ uri="http://localhost:8090";, metalake_name=cls.METALAKE_NAME
+ )
+ cls.temp_dir = tempfile.mkdtemp(prefix="ray_iceberg_it_")
+ cls.gravitino_client.create_catalog(
+ name=cls.CATALOG_NAME,
+ catalog_type=Catalog.Type.RELATIONAL,
+ provider="lakehouse-iceberg",
+ comment="Ray Iceberg IT catalog",
+ # Gravitino requires a non-empty URI property even though its
memory backend
+ # does not use the URI to initialize the in-memory Iceberg catalog.
+ properties={
+ "catalog-backend": "memory",
+ "uri": "memory://ray-iceberg-it",
+ "warehouse": cls.temp_dir,
+ },
+ )
+
+ if not cls._wait_for_iceberg_rest_ready():
+ raise RuntimeError(
+ "Iceberg REST auxiliary service did not become ready at "
+ + ICEBERG_REST_BASE_URL
+ )
+
+ # Imports are deferred so the module can be collected without the
optional IT deps.
+ # pylint: disable=import-error,import-outside-toplevel
+ from pyiceberg.catalog import load_catalog
+ from pyiceberg.schema import Schema
+ from pyiceberg.types import LongType, NestedField, StringType
+
+ # pylint: enable=import-outside-toplevel
+ cls.iceberg_catalog = load_catalog(**cls._catalog_kwargs())
+ cls.iceberg_catalog.create_namespace(cls.SCHEMA_NAME)
+ cls.iceberg_catalog.create_table(
+ cls.TABLE_IDENTIFIER,
+ schema=Schema(
+ NestedField(1, "id", LongType(), required=False),
+ NestedField(2, "value", StringType(), required=False),
+ ),
+ )
+ cls.setup_completed = True
+
+ @classmethod
+ def tearDownClass(cls):
+ try:
+ cls._cleanup_resources()
+ finally:
+ if cls.use_external_gravitino() or cls.gravitino_startup_script is
not None:
+ super().tearDownClass()
+
+ @classmethod
+ def _cleanup_after_setup_failure(cls):
+ if cls.setup_completed:
+ return
+ try:
+ cls._cleanup_resources()
+ finally:
+ if cls.use_external_gravitino() or cls.gravitino_startup_script is
not None:
+ super().tearDownClass()
+
+ @classmethod
+ def _cleanup_resources(cls):
+ failures = []
+
+ try:
+ if cls.iceberg_catalog is not None:
+ cls.iceberg_catalog.drop_table(cls.TABLE_IDENTIFIER)
+ except Exception as error: # pylint: disable=broad-exception-caught
+ failures.append(("drop Iceberg table", error))
+
+ try:
+ if cls.gravitino_client is not None:
+ cls.gravitino_client.drop_catalog(name=cls.CATALOG_NAME,
force=True)
+ except Exception as error: # pylint: disable=broad-exception-caught
+ failures.append(("drop Gravitino catalog", error))
+
+ try:
+ if cls.gravitino_admin_client is not None:
+ cls.gravitino_admin_client.drop_metalake(
+ name=cls.METALAKE_NAME, force=True
+ )
+ except Exception as error: # pylint: disable=broad-exception-caught
+ failures.append(("drop Gravitino metalake", error))
+
+ try:
+ if cls.appended_iceberg_rest_conf and cls.main_conf_path is not
None:
+ cls._reset_conf(
+ {
+ ICEBERG_REST_CONFIG_PROVIDER_KEY:
"dynamic-config-provider",
+ ICEBERG_REST_METALAKE_KEY: cls.METALAKE_NAME,
+ ICEBERG_REST_DEFAULT_CATALOG_KEY: cls.CATALOG_NAME,
+ },
+ cls.main_conf_path,
+ )
+ cls.restart_server()
+ except Exception as error: # pylint: disable=broad-exception-caught
+ failures.append(("reset Iceberg REST conf", error))
+ finally:
+ if cls.temp_dir and os.path.exists(cls.temp_dir):
+ shutil.rmtree(cls.temp_dir, ignore_errors=True)
+
+ for step, error in failures:
+ logger.warning("Cleanup step %s failed: %s", step, error)
+
+ def test_ray_can_write_and_read_existing_iceberg_table(self):
+ # pylint: disable=import-outside-toplevel,import-error
+ import ray
+
+ # pylint: enable=import-outside-toplevel,import-error
+ ray.init(
+ ignore_reinit_error=True,
+ num_cpus=2,
+ include_dashboard=False,
+ log_to_driver=False,
+ )
+ try:
+ catalog_kwargs = self._catalog_kwargs()
+ empty_dataset = ray.data.read_iceberg(
+ table_identifier=self.TABLE_IDENTIFIER,
+ catalog_kwargs=catalog_kwargs,
+ )
+ self.assertEqual(0, empty_dataset.count())
+
+ data = ray.data.from_items(
+ [{"id": index, "value": f"value-{index}"} for index in
range(8)]
+ ).repartition(2)
+ data.write_iceberg(
+ table_identifier=self.TABLE_IDENTIFIER,
+ catalog_kwargs=catalog_kwargs,
+ )
+
+ dataset = ray.data.read_iceberg(
+ table_identifier=self.TABLE_IDENTIFIER,
+ catalog_kwargs=catalog_kwargs,
+ )
+ self.assertEqual(8, dataset.count())
+ self.assertEqual(["id", "value"], dataset.schema().names)
+ self.assertEqual(4, dataset.filter(lambda row: row["id"] % 2 ==
0).count())
+ finally:
+ ray.shutdown()
+
+ @classmethod
+ def _catalog_kwargs(cls):
+ return {
+ "name": "default",
+ "type": "rest",
+ "uri": ICEBERG_REST_BASE_URL,
+ "auth": {
+ "type": "basic",
+ "basic": {"username": "anonymous", "password": ""},
+ },
+ }
+
+ @staticmethod
+ def _wait_for_iceberg_rest_ready(timeout_s: float = 60.0) -> bool:
+ deadline = time.monotonic() + timeout_s
+ while time.monotonic() < deadline:
+ try:
+ response = requests.get(ICEBERG_REST_BASE_URL + "v1/config",
timeout=2)
+ if response.status_code == 200:
+ return True
+ except requests.RequestException:
+ pass
+ time.sleep(0.5)
+ return False
diff --git a/clients/client-python/tests/unittests/test_integration_test_env.py
b/clients/client-python/tests/unittests/test_integration_test_env.py
new file mode 100644
index 0000000000..c2ac040014
--- /dev/null
+++ b/clients/client-python/tests/unittests/test_integration_test_env.py
@@ -0,0 +1,146 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import os
+import shutil
+import subprocess
+import tempfile
+import unittest
+from contextlib import ExitStack
+from pathlib import Path
+from unittest.mock import patch
+
+from gravitino.exceptions.base import GravitinoRuntimeException
+from tests.integration.integration_test_env import IntegrationTestEnv
+
+
+class TestIntegrationTestEnv(unittest.TestCase):
+ def setUp(self):
+ self.patches = ExitStack()
+ self.addCleanup(self.patches.close)
+ home = Path(self.patches.enter_context(tempfile.TemporaryDirectory()))
+ self.script = home / "bin" / "gravitino.sh"
+ self.script.parent.mkdir()
+ self.script.touch()
+ self.data = home / "data"
+ self.data.mkdir()
+ self.metadata = self.data / "metadata"
+ self.metadata.write_text("existing metadata", encoding="utf-8")
+
+ self.patches.enter_context(
+ patch.dict(os.environ, {"GRAVITINO_HOME": str(home)})
+ )
+ self.run = self.patches.enter_context(
+ patch("tests.integration.integration_test_env.subprocess.run")
+ )
+ self.run.return_value = subprocess.CompletedProcess([], 0, "", "")
+ self.ready = self.patches.enter_context(
+ patch(
+
"tests.integration.integration_test_env.check_gravitino_server_status",
+ return_value=True,
+ )
+ )
+ self.client = self.patches.enter_context(
+
patch("tests.integration.integration_test_env.GravitinoAdminClient")
+ )
+ self.patches.enter_context(
+ patch.object(IntegrationTestEnv, "gravitino_admin_client", None)
+ )
+
+ def test_restart_stops_server_before_removing_data(self):
+ events = []
+ remove_data = shutil.rmtree
+
+ def run_command(command, **kwargs):
+ events.append(command[1])
+ self.assertEqual("anonymous", kwargs["env"]["HADOOP_USER_NAME"])
+ return subprocess.CompletedProcess(command, 0, "", "")
+
+ def remove_directory(path):
+ self.assertEqual(["stop"], events)
+ events.append("remove data")
+ remove_data(path)
+
+ self.run.side_effect = run_command
+ self.ready.side_effect = lambda: events.append("ready") or True
+ with patch(
+ "tests.integration.integration_test_env.shutil.rmtree",
+ side_effect=remove_directory,
+ ):
+ IntegrationTestEnv.restart_server()
+
+ self.assertEqual(["stop", "remove data", "start", "ready"], events)
+ self.assertFalse(self.data.exists())
+ self.assertIs(
+ self.client.return_value, IntegrationTestEnv.gravitino_admin_client
+ )
+
+ def test_restart_without_data_directory(self):
+ shutil.rmtree(self.data)
+ with patch("tests.integration.integration_test_env.shutil.rmtree") as
remove:
+ IntegrationTestEnv.restart_server()
+ remove.assert_not_called()
+ self.assertEqual(
+ ["stop", "start"], [call.args[0][1] for call in
self.run.call_args_list]
+ )
+ self.ready.assert_called_once_with()
+
+ def test_failed_stop_preserves_data_and_does_not_start(self):
+ self.run.side_effect = subprocess.CalledProcessError(
+ 1, [str(self.script), "stop"]
+ )
+ with self.assertRaises(subprocess.CalledProcessError):
+ IntegrationTestEnv.restart_server()
+ self.assertEqual("existing metadata",
self.metadata.read_text(encoding="utf-8"))
+ self.assertEqual(1, self.run.call_count)
+ self.ready.assert_not_called()
+ self.client.assert_not_called()
+
+ def test_failed_cleanup_does_not_start(self):
+ with patch(
+ "tests.integration.integration_test_env.shutil.rmtree",
+ side_effect=OSError("Cannot remove test metadata"),
+ ):
+ with self.assertRaises(OSError):
+ IntegrationTestEnv.restart_server()
+ self.assertEqual(1, self.run.call_count)
+ self.assertEqual("stop", self.run.call_args.args[0][1])
+ self.ready.assert_not_called()
+
+ def test_failed_start_does_not_poll_readiness(self):
+ self.run.side_effect = [
+ subprocess.CompletedProcess([str(self.script), "stop"], 0, "", ""),
+ subprocess.CompletedProcess([str(self.script), "start"], 1, "",
"failed"),
+ ]
+ with self.assertRaises(subprocess.CalledProcessError):
+ IntegrationTestEnv.restart_server()
+ self.assertFalse(self.data.exists())
+ self.ready.assert_not_called()
+ self.client.assert_not_called()
+
+ def test_unready_server_does_not_create_client(self):
+ self.ready.return_value = False
+ with self.assertRaises(GravitinoRuntimeException):
+ IntegrationTestEnv.restart_server()
+ self.client.assert_not_called()
+
+ def test_missing_script_preserves_data(self):
+ self.script.unlink()
+ with self.assertRaises(GravitinoRuntimeException):
+ IntegrationTestEnv.restart_server()
+ self.assertTrue(self.metadata.exists())
+ self.run.assert_not_called()
