Merge pull request GoogleCloudPlatform#1399 from Dippatel98:streaming-modes

cloud-teleport · cloud-teleport · commit 7449ae692bac · 2024-03-27T11:51:49.000-07:00
PiperOrigin-RevId: 619603915
diff --git a/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java b/it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/datagenerator/DataGenerator.java
@@ -86,7 +86,7 @@ public static Builder builderWithSchemaTemplate(String testName, String schemaTe
    * @return approximate number of messages generated
    * @throws IOException if any errors are encountered.
    */
-  public Integer execute(Duration timeout) throws IOException {
+  public Long execute(Duration timeout) throws IOException {
     LaunchInfo dataGeneratorLaunchInfo =
         pipelineLauncher.launch(PROJECT, REGION, dataGeneratorOptions);
     assertThatPipeline(dataGeneratorLaunchInfo).isRunning();
@@ -102,11 +102,11 @@ public Integer execute(Duration timeout) throws IOException {
       assertThatResult(dataGeneratorResult).hasTimedOut();
     }
     @SuppressWarnings("nullness")
-    int generatedMessages =
+    Long generatedMessages =
         pipelineLauncher
             .getMetric(
                 PROJECT, REGION, dataGeneratorLaunchInfo.jobId(), MESSAGES_GENERATED_METRIC_NAME)
-            .intValue();
+            .longValue();
     LOG.info("Data generator finished. Generated {} messages.", generatedMessages);
     return generatedMessages;
   }
diff --git a/v1/src/test/java/com/google/cloud/teleport/templates/PubsubToPubsubLT.java b/v1/src/test/java/com/google/cloud/teleport/templates/PubsubToPubsubLT.java
@@ -195,7 +195,7 @@ public void testSteadyState1hr(Function<LaunchConfig.Builder, LaunchConfig.Build
     LaunchInfo info = pipelineLauncher.launch(project, region, options);
     assertThatPipeline(info).isRunning();
     // ElementCount metric in dataflow is approximate, allow for 1% difference
-    Integer expectedMessages = (int) (dataGenerator.execute(Duration.ofMinutes(60)) * 0.99);
+    Long expectedMessages = (long) (dataGenerator.execute(Duration.ofMinutes(60)) * 0.99);
     Result result =
         pipelineOperator.waitForConditionAndCancel(
             createConfig(info, Duration.ofMinutes(20)),
