Skip to content

Commit 661b384

Browse files
committed
Add smoke test for BEAM pipelines
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.
1 parent 81b3a2f commit 661b384

7 files changed

Lines changed: 282 additions & 2 deletions

File tree

core/build.gradle

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -579,6 +579,11 @@ if (environment == 'alpha') {
579579
mainClass: 'google.registry.beam.resave.ResaveAllEppResourcesPipeline',
580580
metaData: 'google/registry/beam/resave_all_epp_resources_pipeline_metadata.json'
581581
],
582+
smokeTest:
583+
[
584+
mainClass: 'google.registry.beam.common.SmokeTestPipeline',
585+
metaData: 'google/registry/beam/smoke_test_pipeline_metadata.json'
586+
],
582587
]
583588
project.tasks.create("stageBeamPipelines") {
584589
doLast {
Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
// Copyright 2026 The Nomulus Authors. All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package google.registry.beam.common;
16+
17+
import static com.google.common.base.Verify.verify;
18+
19+
import com.google.common.flogger.FluentLogger;
20+
import google.registry.model.tld.Tld;
21+
import google.registry.persistence.transaction.CriteriaQueryBuilder;
22+
import java.io.Serializable;
23+
import org.apache.beam.sdk.Pipeline;
24+
import org.apache.beam.sdk.PipelineResult;
25+
import org.apache.beam.sdk.coders.StringUtf8Coder;
26+
import org.apache.beam.sdk.options.PipelineOptionsFactory;
27+
import org.apache.beam.sdk.transforms.Count;
28+
import org.apache.beam.sdk.transforms.DoFn;
29+
import org.apache.beam.sdk.transforms.ParDo;
30+
31+
/**
32+
* For smoke test in the build/deployment process.
33+
*
34+
* <p>There two coverage gaps in unit tests for BEAM pipelines:
35+
*
36+
* <ul>
37+
* <li>The compatibility of the JVM and SDK in the pipeline image
38+
* <li>The JPA setup, which is performed by the {@link RegistryPipelineWorkerInitializer}
39+
* </ul>
40+
*
41+
* <p>This classes defines a pipeline that performs one quick database query. The pipeline is
42+
* expected to complete quickly, and the build or deployment process may launch it on GCP and wait
43+
* for its completion to be certain that all aspects are tested for Nomulus pipelines.
44+
*/
45+
public class SmokeTestPipeline implements Serializable {
46+
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
47+
48+
public static void main(String[] args) {
49+
PipelineOptionsFactory.register(RegistryPipelineOptions.class);
50+
RegistryPipelineOptions options =
51+
PipelineOptionsFactory.fromArgs(args).withValidation().as(RegistryPipelineOptions.class);
52+
runPipeline(options);
53+
}
54+
55+
static PipelineResult runPipeline(RegistryPipelineOptions options) {
56+
Pipeline pipeline = Pipeline.create(options);
57+
pipeline
58+
.apply(
59+
"Read Tlds",
60+
RegistryJpaIO.read(() -> CriteriaQueryBuilder.create(Tld.class).build(), Tld::getTldStr)
61+
.withCoder(StringUtf8Coder.of()))
62+
.apply("Count Tlds", Count.globally())
63+
.apply(
64+
"Verify Count",
65+
ParDo.of(
66+
new DoFn<Long, Void>() {
67+
@DoFn.ProcessElement
68+
public void processElement(@Element Long count) {
69+
logger.atInfo().log("Tld count: %s", count);
70+
verify(count > 0, "Expecting 1 or more, got %s.", count);
71+
}
72+
}));
73+
74+
return pipeline.run();
75+
}
76+
}
Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
{
2+
"name": "Beam pipeline smoke test",
3+
"description": "An Apache Beam pipeline that performs a simple database query.",
4+
"parameters": [
5+
{
6+
"name": "registryEnvironment",
7+
"label": "The Registry environment.",
8+
"helpText": "The Registry environment.",
9+
"is_optional": false,
10+
"regexes": [
11+
"^SANDBOX|CRASH$"
12+
]
13+
},
14+
{
15+
"name": "isolationOverride",
16+
"label": "The desired SQL transaction isolation level.",
17+
"helpText": "The desired SQL transaction isolation level.",
18+
"is_optional": true,
19+
"regexes": [
20+
"^[0-9A-Z_]+$"
21+
]
22+
}
23+
]
24+
}
Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
// Copyright 2026 The Nomulus Authors. All Rights Reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package google.registry.beam.common;
16+
17+
import static com.google.common.truth.Truth.assertThat;
18+
import static google.registry.persistence.PersistenceModule.TransactionIsolationLevel.TRANSACTION_REPEATABLE_READ;
19+
import static google.registry.testing.DatabaseHelper.createTld;
20+
import static org.junit.jupiter.api.Assertions.assertThrows;
21+
22+
import google.registry.beam.TestPipelineExtension;
23+
import google.registry.persistence.transaction.JpaTestExtensions;
24+
import google.registry.testing.FakeClock;
25+
import java.time.Instant;
26+
import org.apache.beam.sdk.Pipeline;
27+
import org.apache.beam.sdk.options.PipelineOptionsFactory;
28+
import org.hibernate.cfg.AvailableSettings;
29+
import org.junit.jupiter.api.Test;
30+
import org.junit.jupiter.api.extension.RegisterExtension;
31+
32+
public class SmokeTestPipelineTest {
33+
34+
private final FakeClock clock = new FakeClock(Instant.parse("2021-02-02T00:00:05.000Z"));
35+
36+
@RegisterExtension
37+
final JpaTestExtensions.JpaIntegrationTestExtension jpa =
38+
new JpaTestExtensions.Builder()
39+
.withClock(clock)
40+
.withProperty(AvailableSettings.ISOLATION, TRANSACTION_REPEATABLE_READ.name())
41+
.buildIntegrationTestExtension();
42+
43+
@RegisterExtension
44+
final TestPipelineExtension pipeline =
45+
TestPipelineExtension.create().enableAbandonedNodeEnforcement(true);
46+
47+
private final RegistryPipelineOptions options =
48+
PipelineOptionsFactory.create().as(RegistryPipelineOptions.class);
49+
50+
@Test
51+
void whenIldsDoNotExist_failure() {
52+
var exception =
53+
assertThrows(
54+
Pipeline.PipelineExecutionException.class,
55+
() -> SmokeTestPipeline.runPipeline(options).waitUntilFinish());
56+
assertThat(exception).hasMessageThat().contains("Expecting 1 or more, got 0.");
57+
}
58+
59+
@Test
60+
void whenTldsExist_success() {
61+
createTld("tld");
62+
SmokeTestPipeline.runPipeline(options).waitUntilFinish();
63+
}
64+
}

release/cloudbuild-nomulus.yaml

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,9 @@ steps:
157157
google.registry.beam.rde.RdePipeline \
158158
google/registry/beam/rde_pipeline_metadata.json \
159159
google.registry.beam.resave.ResaveAllEppResourcesPipeline \
160-
google/registry/beam/resave_all_epp_resources_pipeline_metadata.json
160+
google/registry/beam/resave_all_epp_resources_pipeline_metadata.json \
161+
google.registry.beam.common.SmokeTestPipeline \
162+
google/registry/beam/smoke_test_pipeline_metadata.json
161163
# Build and upload the schema jar as well as other artifacts needed by the schema tests.
162164
- name: 'gcr.io/${PROJECT_ID}/builder:latest'
163165
entrypoint: /bin/bash

release/cloudbuild-release.yaml

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# To manually trigger a build on GCB, run:
22
# gcloud builds submit --config cloudbuild-release.yaml --substitutions \
3-
# TAG_NAME=[TAG],_INTERNAL_REPO_URL=[URL] ..
3+
# TAG_NAME=[TAG],_INTERNAL_REPO_URL=[URL],_TEST_PROJECT=[_TEST_PROJECT] \
4+
# ..
45
#
56
# To trigger a build automatically, follow the instructions below and add a trigger:
67
# https://cloud.google.com/cloud-build/docs/running-builds/automate-builds
@@ -288,6 +289,19 @@ steps:
288289
echo "Tag format '$TAG_NAME' does not match a known release type. Exiting."
289290
exit 1
290291
fi
292+
# Run the BEAM smoke test, using the builder and pipeline image just created
293+
- name: 'gcr.io/$PROJECT_ID/builder:latest'
294+
entrypoint: /bin/bash
295+
args:
296+
- -c
297+
- |
298+
set -e
299+
if [[ "${TAG_NAME}" =~ ^nomulus-20[0-9]{2}[0-1][0-9][0-3][0-9]-RC[0-9]{2}$ ]]; then
300+
gcloud secrets versions access latest \
301+
--secret nomulus-tool-cloudbuild-credential > tool-credential.json
302+
gcloud auth activate-service-account --key-file=tool-credential.json
303+
./release/run_beam_smoketest.sh "${TAG_NAME}" "${PROJECT_ID}" "${_TEST_PROJECT}"
304+
fi
291305
timeout: 3600s
292306
options:
293307
machineType: 'E2_HIGHCPU_32'

release/run_beam_smoketest.sh

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
#!/bin/bash
2+
# Copyright 2026 The Nomulus Authors. All Rights Reserved.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
#
16+
# This script runs the BEAM pipeline smoke test as part of the build process.
17+
# It assumes that all pipelines have been built and staged.
18+
#
19+
# This script expects the following arguments in order:
20+
# - The release tag
21+
# - The GCP project that serves the release artifacts
22+
# - The GCP project id where the smoke test should run. This id should have the
23+
# Nomulus environment name as suffix, following the last '-' in the project id,
24+
# E.g., domain-registry-crash. Only crash and sandbox are allowed as the test
25+
# environment. Use 'NONE' as project id to skip this test.
26+
# - GCP region to run the test. This is optional.
27+
28+
set -e
29+
30+
if [[ $# -ne 3 && $# -ne 4 ]];
31+
then
32+
echo "Usage: $0 <release_tag> <build_project_id> <test_project_id> [<gcp_region>]"
33+
exit 1
34+
fi
35+
36+
release_tag="$1"
37+
build_project="$2"
38+
test_project="$3"
39+
region="${4:-us-central1}"
40+
41+
if [[ "${test_project}" == "NONE" ]]; then
42+
echo "BEAM smoke test skipped as requested."
43+
exit 0
44+
fi
45+
46+
test_project_id_suffix="${test_project##*-}"
47+
if [[ "${test_project}" == *"-"* ]]; then
48+
# Convert environment name to upper case
49+
test_env="${test_project_id_suffix^^}"
50+
else
51+
test_env=""
52+
fi
53+
54+
if [[ -z "${test_env}" ]]; then
55+
echo "Cannot extract environment from project id ${test_project}"
56+
exit 1
57+
fi
58+
59+
if [[ ${test_env} != "CRASH" && ${test_env} != "SANDBOX" ]]; then
60+
echo "Expecting CRASH or SANDBOX as test environment, got ${test_env}"
61+
exit 1
62+
fi
63+
64+
template_folder="gs://${build_project}-deploy/${release_tag}/beam"
65+
template="${template_folder}/smoke_test_pipeline_metadata.json"
66+
job_name=$(echo "beam-smoketest-${release_tag}" | tr '[:upper:]_' '[:lower:]-')
67+
job_id=$(gcloud dataflow flex-template run "${job_name}" \
68+
--template-file-gcs-location="${template}" \
69+
--region="${region}" --parameters=registryEnvironment="${test_env}" \
70+
--project=${test_project} --format='value(job.id)')
71+
echo "Test pipeline started as ${job_id}"
72+
73+
# Wait up to 30 minutes for the smoke test to finish. This 5X of a typical
74+
# run.
75+
for i in {1..30}
76+
do
77+
job_state=$(gcloud dataflow jobs describe "${job_id}" --region=${region} \
78+
--format="value(currentState)" --project "${test_project}")
79+
echo "Test pipeline state is ${job_state}"
80+
81+
if [[ "${job_state}" == "JOB_STATE_DONE" ]]; then
82+
echo "Smoke test completed successfully."
83+
exit 0
84+
elif [[ "${job_state}" == "JOB_STATE_QUEUED" || \
85+
"${job_state}" == "JOB_STATE_RUNNING" ]]; then
86+
echo "Sleeping for 60 seconds"
87+
sleep 60
88+
else
89+
echo "Unexpected job state ${job_state}"
90+
exit 1
91+
fi
92+
done
93+
94+
echo "Error: Smoke test did not complete in time."
95+
exit 1

0 commit comments

Comments
 (0)