From 94adacdcfb1a07cb21a843eabd1c09f56061cde8 Mon Sep 17 00:00:00 2001
From: Valentyn Tymofieiev <valentyn@google.com>
Date: Wed, 20 Sep 2023 14:48:39 -0700
Subject: [PATCH] Use a single marker for Vertex AI tests to not run them
 twice.

---
 .../apache_beam/ml/inference/vertex_ai_inference_it_test.py    | 3 +--
 sdks/python/pytest.ini                                         | 2 +-
 sdks/python/test-suites/dataflow/common.gradle                 | 2 +-
 3 files changed, 3 insertions(+), 4 deletions(-)

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"
       ]