diff --git a/sdks/python/apache_beam/ml/inference/vertex_ai_inference_it_test.py b/sdks/python/apache_beam/ml/inference/vertex_ai_inference_it_test.py index 8ba315ba00ca..168ab031abb1 100644 --- a/sdks/python/apache_beam/ml/inference/vertex_ai_inference_it_test.py +++ b/sdks/python/apache_beam/ml/inference/vertex_ai_inference_it_test.py @@ -46,8 +46,7 @@ class VertexAIInference(unittest.TestCase): - @pytest.mark.uses_vertex_ai - @pytest.mark.it_postcommit + @pytest.mark.vertex_ai_postcommit def test_vertex_ai_run_flower_image_classification(self): output_file = '/'.join([_OUTPUT_DIR, str(uuid.uuid4()), 'output.txt']) diff --git a/sdks/python/pytest.ini b/sdks/python/pytest.ini index ec8d9fb3900b..d1d69f798355 100644 --- a/sdks/python/pytest.ini +++ b/sdks/python/pytest.ini @@ -60,7 +60,7 @@ markers = uses_onnx: tests that utilizes onnx in some way. uses_tf: tests that utilize tensorflow. uses_transformers: tests that utilize transformers in some way. - uses_vertex_ai: tests that utilize vertex ai in some way. + vertex_ai_postcommit: vertex ai postcommits that need additional deps. # Default timeout intended for unit tests. # If certain tests need a different value, please see the docs on how to diff --git a/sdks/python/test-suites/dataflow/common.gradle b/sdks/python/test-suites/dataflow/common.gradle index 8b8a56808996..f907a728c0b5 100644 --- a/sdks/python/test-suites/dataflow/common.gradle +++ b/sdks/python/test-suites/dataflow/common.gradle @@ -417,7 +417,7 @@ task vertexAIInferenceTest { def argMap = [ "test_opts": testOpts, "suite": "VertexAITests-df-py${pythonVersionSuffix}", - "collect": "uses_vertex_ai and it_postcommit" , + "collect": "vertex_ai_postcommit" , "runner": "TestDataflowRunner", "requirements_file": "$requirementsFile" ]