fix unit tests

Author: bambadiouf1

runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslatorTest.java

@@ -620,31 +620,54 @@ public void testDiskSizeGbConfig() throws IOException {
   }
 
   @Test
-  public void testDiskProvisioningTranslation() {
-    DataflowPipelineOptions options =
-        PipelineOptionsFactory.create().as(DataflowPipelineOptions.class);
+  public void testDiskProvisioningTranslation() throws IOException {
+    DataflowPipelineOptions options = buildPipelineOptions();
     options.setDiskProvisionedIops(Long.valueOf(7000));
     options.setDiskProvisionedThroughputMibps(Long.valueOf(250));
-    options.setProject("test-project"); // Required for translator
 
-    WorkerPool pool = translateWorkerPool(options);
+    Pipeline p = buildPipeline(options);
+    p.traverseTopologically(new RecordingPipelineVisitor());
+    SdkComponents sdkComponents = createSdkComponents(options);
+    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p, sdkComponents, true);
+    Job job =
+        DataflowPipelineTranslator.fromOptions(options)
+            .translate(
+                p,
+                pipelineProto,
+                sdkComponents,
+                DataflowRunner.fromOptions(options),
+                Collections.emptyList())
+            .getJob();
 
+    assertEquals(1, job.getEnvironment().getWorkerPools().size());
+    WorkerPool pool = job.getEnvironment().getWorkerPools().get(0);
     assertEquals(Long.valueOf(7000), pool.getDiskProvisionedIops());
     assertEquals(Long.valueOf(250), pool.getDiskProvisionedThroughputMibps());
   }
 
   @Test
-  public void testDiskProvisioningTranslationDefaults() {
-    DataflowPipelineOptions options =
-        PipelineOptionsFactory.create().as(DataflowPipelineOptions.class);
-    options.setProject("test-project"); // Required for translator
+  public void testDiskProvisioningTranslationDefaults() throws IOException {
+    DataflowPipelineOptions options = buildPipelineOptions();
 
-    WorkerPool pool = translateWorkerPool(options);
+    Pipeline p = buildPipeline(options);
+    p.traverseTopologically(new RecordingPipelineVisitor());
+    SdkComponents sdkComponents = createSdkComponents(options);
+    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p, sdkComponents, true);
+    Job job =
+        DataflowPipelineTranslator.fromOptions(options)
+            .translate(
+                p,
+                pipelineProto,
+                sdkComponents,
+                DataflowRunner.fromOptions(options),
+                Collections.emptyList())
+            .getJob();
 
+    assertEquals(1, job.getEnvironment().getWorkerPools().size());
+    WorkerPool pool = job.getEnvironment().getWorkerPools().get(0);
     assertNull(pool.getDiskProvisionedIops());
     assertNull(pool.getDiskProvisionedThroughputMibps());
   }
-
   /** A composite transform that returns an output that is unrelated to the input. */
   private static class UnrelatedOutputCreator
       extends PTransform<PCollection<Integer>, PCollection<Integer>> {