Be able to pass success to BigQueryOperation

contrib/flo-bigquery/src/main/java/com/spotify/flo/contrib/bigquery/BigQueryOperation.java

@@ -61,6 +61,10 @@ static <T> BigQueryOperation<T> ofJob(Fn<JobInfo> job) {
     return new BigQueryOperation<T>().job(job);
   }
 
+  static <T> BigQueryOperation<T> ofJob(Fn<JobInfo> job, F1<JobInfo, T> success) {
+    return new BigQueryOperation<T>().job(job).success(success);
+  }
+
   public static class Provider<T> implements Serializable {
 
     private static final long serialVersionUID = 1L;
@@ -95,5 +99,9 @@ public BigQueryOperation<T> job(Fn<JobInfo> jobInfo) {
     public BigQueryOperation<T> job(JobInfo jobInfo) {
       return BigQueryOperation.ofJob(() -> jobInfo);
     }
+
+    public BigQueryOperation<T> job(JobInfo jobInfo, F1<JobInfo, T> success) {
+      return BigQueryOperation.ofJob(() -> jobInfo, success);
+    }
   }
 }

contrib/flo-bigquery/src/test/java/com/spotify/flo/contrib/bigquery/BigQueryOperatorTest.java

@@ -120,6 +120,26 @@ public void shouldRunLoadJobInTestMode() throws Exception {
     }
   }
 
+  @Test
+  public void shouldRunLoadJobInTestModeWithSuccess() throws Exception {
+    final TableId dstTable = TableId.of("foo", "bar", "baz");
+    final String srcUri = "gs://foo/bar";
+
+    final Task<TableId> task = Task.named("task")
+        .ofType(TableId.class)
+        .operator(BigQueryOperator.create())
+        .process(bq -> bq.job(
+            JobInfo.of(LoadJobConfiguration.of(dstTable, srcUri)), response -> dstTable));
+
+    try (TestScope scope = FloTesting.scope()) {
+
+      final TableId result = FloRunner.runTask(task).future()
+          .get(30, SECONDS);
+
+      assertThat(result, is(dstTable));
+    }
+  }
+
   @Test
   public void shouldRunExtractJobInTestMode() throws Exception {
     final TableId srcTable = TableId.of("foo", "bar", "baz");