From 60d3653b46d639fc6ff7c173b0d93de70a18e499 Mon Sep 17 00:00:00 2001 From: Weimin Yu Date: Thu, 7 May 2026 16:41:13 -0400 Subject: [PATCH] Add smoke test for BEAM pipelines (#3037) Created a smoke test to cover unit test gaps wrt BEAM: - The Java and SDK compatibility in the pipeline container image - The JPA setup in the pipelines Both issues above can only be tested in a real pipeline. This PR defines a new pipeline that performs a lightweight SQL query and minimal processing. The build process can launch it in a test environment to verify that the pipelines in the build can run. The run script is also provided. --- core/build.gradle | 5 + .../beam/common/SmokeTestPipeline.java | 76 +++++++++++++++ .../beam/smoke_test_pipeline_metadata.json | 24 +++++ .../beam/common/SmokeTestPipelineTest.java | 64 +++++++++++++ release/cloudbuild-nomulus.yaml | 4 +- release/cloudbuild-release.yaml | 18 +++- release/run_beam_smoketest.sh | 96 +++++++++++++++++++ 7 files changed, 284 insertions(+), 3 deletions(-) create mode 100644 core/src/main/java/google/registry/beam/common/SmokeTestPipeline.java create mode 100644 core/src/main/resources/google/registry/beam/smoke_test_pipeline_metadata.json create mode 100644 core/src/test/java/google/registry/beam/common/SmokeTestPipelineTest.java create mode 100755 release/run_beam_smoketest.sh diff --git a/core/build.gradle b/core/build.gradle index 12b127b81..2f288fb48 100644 --- a/core/build.gradle +++ b/core/build.gradle @@ -579,6 +579,11 @@ if (environment == 'alpha') { mainClass: 'google.registry.beam.resave.ResaveAllEppResourcesPipeline', metaData: 'google/registry/beam/resave_all_epp_resources_pipeline_metadata.json' ], + smokeTest: + [ + mainClass: 'google.registry.beam.common.SmokeTestPipeline', + metaData: 'google/registry/beam/smoke_test_pipeline_metadata.json' + ], ] project.tasks.create("stageBeamPipelines") { doLast { diff --git a/core/src/main/java/google/registry/beam/common/SmokeTestPipeline.java b/core/src/main/java/google/registry/beam/common/SmokeTestPipeline.java new file mode 100644 index 000000000..093e9eb84 --- /dev/null +++ b/core/src/main/java/google/registry/beam/common/SmokeTestPipeline.java @@ -0,0 +1,76 @@ +// Copyright 2026 The Nomulus Authors. All Rights Reserved. +// +// Licensed 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. + +package google.registry.beam.common; + +import static com.google.common.base.Verify.verify; + +import com.google.common.flogger.FluentLogger; +import google.registry.model.tld.Tld; +import google.registry.persistence.transaction.CriteriaQueryBuilder; +import java.io.Serializable; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.transforms.Count; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; + +/** + * For smoke test in the build/deployment process. + * + *

There two coverage gaps in unit tests for BEAM pipelines: + * + *

+ * + *

This classes defines a pipeline that performs one quick database query. The pipeline is + * expected to complete quickly, and the build or deployment process may launch it on GCP and wait + * for its completion to be certain that all aspects are tested for Nomulus pipelines. + */ +public class SmokeTestPipeline implements Serializable { + private static final FluentLogger logger = FluentLogger.forEnclosingClass(); + + public static void main(String[] args) { + PipelineOptionsFactory.register(RegistryPipelineOptions.class); + RegistryPipelineOptions options = + PipelineOptionsFactory.fromArgs(args).withValidation().as(RegistryPipelineOptions.class); + runPipeline(options); + } + + static PipelineResult runPipeline(RegistryPipelineOptions options) { + Pipeline pipeline = Pipeline.create(options); + pipeline + .apply( + "Read Tlds", + RegistryJpaIO.read(() -> CriteriaQueryBuilder.create(Tld.class).build(), Tld::getTldStr) + .withCoder(StringUtf8Coder.of())) + .apply("Count Tlds", Count.globally()) + .apply( + "Verify Count", + ParDo.of( + new DoFn() { + @DoFn.ProcessElement + public void processElement(@Element Long count) { + logger.atInfo().log("Tld count: %s", count); + verify(count > 0, "Expecting 1 or more, got %s.", count); + } + })); + + return pipeline.run(); + } +} diff --git a/core/src/main/resources/google/registry/beam/smoke_test_pipeline_metadata.json b/core/src/main/resources/google/registry/beam/smoke_test_pipeline_metadata.json new file mode 100644 index 000000000..18fe1f607 --- /dev/null +++ b/core/src/main/resources/google/registry/beam/smoke_test_pipeline_metadata.json @@ -0,0 +1,24 @@ +{ + "name": "Beam pipeline smoke test", + "description": "An Apache Beam pipeline that performs a simple database query.", + "parameters": [ + { + "name": "registryEnvironment", + "label": "The Registry environment.", + "helpText": "The Registry environment.", + "is_optional": false, + "regexes": [ + "^SANDBOX|CRASH$" + ] + }, + { + "name": "isolationOverride", + "label": "The desired SQL transaction isolation level.", + "helpText": "The desired SQL transaction isolation level.", + "is_optional": true, + "regexes": [ + "^[0-9A-Z_]+$" + ] + } + ] +} diff --git a/core/src/test/java/google/registry/beam/common/SmokeTestPipelineTest.java b/core/src/test/java/google/registry/beam/common/SmokeTestPipelineTest.java new file mode 100644 index 000000000..a46787dda --- /dev/null +++ b/core/src/test/java/google/registry/beam/common/SmokeTestPipelineTest.java @@ -0,0 +1,64 @@ +// Copyright 2026 The Nomulus Authors. All Rights Reserved. +// +// Licensed 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. + +package google.registry.beam.common; + +import static com.google.common.truth.Truth.assertThat; +import static google.registry.persistence.PersistenceModule.TransactionIsolationLevel.TRANSACTION_REPEATABLE_READ; +import static google.registry.testing.DatabaseHelper.createTld; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import google.registry.beam.TestPipelineExtension; +import google.registry.persistence.transaction.JpaTestExtensions; +import google.registry.testing.FakeClock; +import java.time.Instant; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.hibernate.cfg.AvailableSettings; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +public class SmokeTestPipelineTest { + + private final FakeClock clock = new FakeClock(Instant.parse("2021-02-02T00:00:05.000Z")); + + @RegisterExtension + final JpaTestExtensions.JpaIntegrationTestExtension jpa = + new JpaTestExtensions.Builder() + .withClock(clock) + .withProperty(AvailableSettings.ISOLATION, TRANSACTION_REPEATABLE_READ.name()) + .buildIntegrationTestExtension(); + + @RegisterExtension + final TestPipelineExtension pipeline = + TestPipelineExtension.create().enableAbandonedNodeEnforcement(true); + + private final RegistryPipelineOptions options = + PipelineOptionsFactory.create().as(RegistryPipelineOptions.class); + + @Test + void whenIldsDoNotExist_failure() { + var exception = + assertThrows( + Pipeline.PipelineExecutionException.class, + () -> SmokeTestPipeline.runPipeline(options).waitUntilFinish()); + assertThat(exception).hasMessageThat().contains("Expecting 1 or more, got 0."); + } + + @Test + void whenTldsExist_success() { + createTld("tld"); + SmokeTestPipeline.runPipeline(options).waitUntilFinish(); + } +} diff --git a/release/cloudbuild-nomulus.yaml b/release/cloudbuild-nomulus.yaml index 703d6f3a6..113eef680 100644 --- a/release/cloudbuild-nomulus.yaml +++ b/release/cloudbuild-nomulus.yaml @@ -157,7 +157,9 @@ steps: google.registry.beam.rde.RdePipeline \ google/registry/beam/rde_pipeline_metadata.json \ google.registry.beam.resave.ResaveAllEppResourcesPipeline \ - google/registry/beam/resave_all_epp_resources_pipeline_metadata.json + google/registry/beam/resave_all_epp_resources_pipeline_metadata.json \ + google.registry.beam.common.SmokeTestPipeline \ + google/registry/beam/smoke_test_pipeline_metadata.json # Build and upload the schema jar as well as other artifacts needed by the schema tests. - name: 'gcr.io/${PROJECT_ID}/builder:latest' entrypoint: /bin/bash diff --git a/release/cloudbuild-release.yaml b/release/cloudbuild-release.yaml index f02555fdd..f23a0d7f9 100644 --- a/release/cloudbuild-release.yaml +++ b/release/cloudbuild-release.yaml @@ -1,6 +1,7 @@ # To manually trigger a build on GCB, run: # gcloud builds submit --config cloudbuild-release.yaml --substitutions \ -# TAG_NAME=[TAG],_INTERNAL_REPO_URL=[URL] .. +# TAG_NAME=[TAG],_INTERNAL_REPO_URL=[URL],_TEST_PROJECT=[_TEST_PROJECT] \ +# .. # # To trigger a build automatically, follow the instructions below and add a trigger: # https://cloud.google.com/cloud-build/docs/running-builds/automate-builds @@ -288,6 +289,19 @@ steps: echo "Tag format '$TAG_NAME' does not match a known release type. Exiting." exit 1 fi -timeout: 3600s +# Run the BEAM smoke test, using the builder and pipeline image just created +- name: 'gcr.io/$PROJECT_ID/builder:latest' + entrypoint: /bin/bash + args: + - -c + - | + set -e + if [[ "${TAG_NAME}" =~ ^nomulus-20[0-9]{2}[0-1][0-9][0-3][0-9]-RC[0-9]{2}$ ]]; then + gcloud secrets versions access latest \ + --secret nomulus-tool-cloudbuild-credential > tool-credential.json + gcloud auth activate-service-account --key-file=tool-credential.json + ./release/run_beam_smoketest.sh "${TAG_NAME}" "${PROJECT_ID}" "${_TEST_PROJECT}" + fi +timeout: 5400s options: machineType: 'E2_HIGHCPU_32' diff --git a/release/run_beam_smoketest.sh b/release/run_beam_smoketest.sh new file mode 100755 index 000000000..ca0ddf2af --- /dev/null +++ b/release/run_beam_smoketest.sh @@ -0,0 +1,96 @@ +#!/bin/bash +# Copyright 2026 The Nomulus Authors. All Rights Reserved. +# +# Licensed 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. +# +# This script runs the BEAM pipeline smoke test as part of the build process. +# It assumes that all pipelines have been built and staged. +# +# This script expects the following arguments in order: +# - The release tag +# - The GCP project that serves the release artifacts +# - The GCP project id where the smoke test should run. This id should have the +# Nomulus environment name as suffix, following the last '-' in the project id, +# E.g., domain-registry-crash. Only crash and sandbox are allowed as the test +# environment. Use 'NONE' as project id to skip this test. +# - GCP region to run the test. This is optional. + +set -e + +if [[ $# -ne 3 && $# -ne 4 ]]; +then + echo "Usage: $0 []" + exit 1 +fi + +release_tag="$1" +build_project="$2" +test_project="$3" +region="${4:-us-central1}" + +if [[ "${test_project}" == "NONE" ]]; then + echo "BEAM smoke test skipped as requested." + exit 0 +fi + +test_project_id_suffix="${test_project##*-}" +if [[ "${test_project}" == *"-"* ]]; then + # Convert environment name to upper case + test_env="${test_project_id_suffix^^}" +else + test_env="" +fi + +if [[ -z "${test_env}" ]]; then + echo "Cannot extract environment from project id ${test_project}" + exit 1 +fi + +if [[ ${test_env} != "CRASH" && ${test_env} != "SANDBOX" ]]; then + echo "Expecting CRASH or SANDBOX as test environment, got ${test_env}" + exit 1 +fi + +template_folder="gs://${build_project}-deploy/${release_tag}/beam" +template="${template_folder}/smoke_test_pipeline_metadata.json" +job_name=$(echo "beam-smoketest-${release_tag}" | tr '[:upper:]_' '[:lower:]-') +job_id=$(gcloud dataflow flex-template run "${job_name}" \ + --template-file-gcs-location="${template}" \ + --region="${region}" --parameters=registryEnvironment="${test_env}" \ + --project=${test_project} --format='value(job.id)') +echo "Test pipeline started as ${job_id}" + +# Wait up to 30 minutes for the smoke test to finish. This 5X of a typical +# run. +for i in {1..30} +do + job_state=$(gcloud dataflow jobs describe "${job_id}" --region=${region} \ + --format="value(currentState)" --project "${test_project}") + echo "Test pipeline state is ${job_state}" + + if [[ "${job_state}" == "JOB_STATE_DONE" ]]; then + echo "Smoke test completed successfully." + exit 0 + elif [[ "${job_state}" == "JOB_STATE_QUEUED" || \ + "${job_state}" == "JOB_STATE_RUNNING" || \ + "${job_state}" == "JOB_STATE_PENDING" ]]; then + echo "Sleeping for 60 seconds" + sleep 60 + else + echo "Unexpected job state ${job_state}" + exit 1 + fi +done + +echo "Error: Smoke test did not complete in time." +exit 1