From a1ec0552a340aa99a6d0aadc77f49caf8acab110 Mon Sep 17 00:00:00 2001
From: Jack McCluskey <34928439+jrmccluskey@users.noreply.github.com>
Date: Tue, 19 Sep 2023 18:18:12 -0400
Subject: [PATCH] Add extra marks to VertexAI IT tests (#28534)

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

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 4ef42fb10a70..8ba315ba00ca 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
@@ -26,6 +26,7 @@
 from apache_beam.io.filesystems import FileSystems
 from apache_beam.testing.test_pipeline import TestPipeline
 
+# pylint: disable=ungrouped-imports
 try:
   from apache_beam.examples.inference import vertex_ai_image_classification
   from apache_beam.examples.inference import vertex_ai_llm_text_classification
diff --git a/sdks/python/test-suites/dataflow/common.gradle b/sdks/python/test-suites/dataflow/common.gradle
index eee6e9d21887..8b8a56808996 100644
--- a/sdks/python/test-suites/dataflow/common.gradle
+++ b/sdks/python/test-suites/dataflow/common.gradle
@@ -418,7 +418,8 @@ task vertexAIInferenceTest {
           "test_opts": testOpts,
           "suite": "VertexAITests-df-py${pythonVersionSuffix}",
           "collect": "uses_vertex_ai and it_postcommit" ,
-          "runner": "TestDataflowRunner"
+          "runner": "TestDataflowRunner",
+          "requirements_file": "$requirementsFile"
       ]
       def cmdArgs = mapToArgString(argMap)
       exec {