add tests

Author: david-leifker

datahub-upgrade/build.gradle

@@ -88,6 +88,9 @@ dependencies {
   testImplementation 'uk.org.webcompere:system-stubs-testng:2.1.7'
   testRuntimeOnly externalDependency.logbackClassic
 
+  testImplementation externalDependency.h2
+  testImplementation testFixtures(project(':metadata-io'))
+
   constraints {
     implementation(implementation externalDependency.parquetHadoop) {
       because("CVE-2022-42003")

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

@@ -2,6 +2,7 @@
 
 import static com.linkedin.metadata.Constants.ASPECT_LATEST_VERSION;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.linkedin.datahub.upgrade.UpgradeContext;
 import com.linkedin.datahub.upgrade.UpgradeStep;
 import com.linkedin.datahub.upgrade.UpgradeStepResult;
@@ -80,9 +81,11 @@ private List<RestoreIndicesResult> iterateFutures(List<Future<RestoreIndicesResu
       if (future.isDone()) {
         try {
           result.add(future.get());
-          futures.remove(future);
         } catch (InterruptedException | ExecutionException e) {
           log.error("Error iterating futures", e);
+          result.add(null); // add null to indicate failure
+        } finally {
+          futures.remove(future);
         }
       }
     }
@@ -152,7 +155,8 @@ private RestoreIndicesArgs getArgs(UpgradeContext context) {
     return result;
   }
 
-  private int getRowCount(RestoreIndicesArgs args) {
+  @VisibleForTesting
+  int getRowCount(RestoreIndicesArgs args) {
     ExpressionList<EbeanAspectV2> countExp =
         _server
             .find(EbeanAspectV2.class)
@@ -234,6 +238,10 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
         while (futures.size() > 0) {
           List<RestoreIndicesResult> tmpResults = iterateFutures(futures);
           for (RestoreIndicesResult tmpResult : tmpResults) {
+            if (tmpResult == null) {
+              // return error if there was an error processing a future
+              return new DefaultUpgradeStepResult(id(), DataHubUpgradeState.FAILED);
+            }
             reportStats(context, finalJobResult, tmpResult, rowCount, startTime);
           }
         }