Skip to content

Commit 0acbe89

Browse files
committed
add tests
1 parent f2e0fd9 commit 0acbe89

File tree

5 files changed

+406
-2
lines changed

5 files changed

+406
-2
lines changed

datahub-upgrade/build.gradle

+3
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,9 @@ dependencies {
8888
testImplementation 'uk.org.webcompere:system-stubs-testng:2.1.7'
8989
testRuntimeOnly externalDependency.logbackClassic
9090

91+
testImplementation externalDependency.h2
92+
testImplementation testFixtures(project(':metadata-io'))
93+
9194
constraints {
9295
implementation(implementation externalDependency.parquetHadoop) {
9396
because("CVE-2022-42003")

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java

+10-2
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION;
44

5+
import com.google.common.annotations.VisibleForTesting;
56
import com.linkedin.datahub.upgrade.UpgradeContext;
67
import com.linkedin.datahub.upgrade.UpgradeStep;
78
import com.linkedin.datahub.upgrade.UpgradeStepResult;
@@ -80,9 +81,11 @@ private List<RestoreIndicesResult> iterateFutures(List<Future<RestoreIndicesResu
8081
if (future.isDone()) {
8182
try {
8283
result.add(future.get());
83-
futures.remove(future);
8484
} catch (InterruptedException | ExecutionException e) {
8585
log.error("Error iterating futures", e);
86+
result.add(null); // add null to indicate failure
87+
} finally {
88+
futures.remove(future);
8689
}
8790
}
8891
}
@@ -152,7 +155,8 @@ private RestoreIndicesArgs getArgs(UpgradeContext context) {
152155
return result;
153156
}
154157

155-
private int getRowCount(RestoreIndicesArgs args) {
158+
@VisibleForTesting
159+
int getRowCount(RestoreIndicesArgs args) {
156160
ExpressionList<EbeanAspectV2> countExp =
157161
_server
158162
.find(EbeanAspectV2.class)
@@ -234,6 +238,10 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
234238
while (futures.size() > 0) {
235239
List<RestoreIndicesResult> tmpResults = iterateFutures(futures);
236240
for (RestoreIndicesResult tmpResult : tmpResults) {
241+
if (tmpResult == null) {
242+
// return error if there was an error processing a future
243+
return new DefaultUpgradeStepResult(id(), DataHubUpgradeState.FAILED);
244+
}
237245
reportStats(context, finalJobResult, tmpResult, rowCount, startTime);
238246
}
239247
}

0 commit comments

Comments
 (0)