Add extra marks to VertexAI IT tests

apache · Sep 19, 2023 · f2b92ec · f2b92ec
1 parent 8f60924
commit f2b92ec
Showing 1 changed file with 5 additions and 0 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
@@ -27,6 +27,7 @@
 from apache_beam.testing.test_pipeline import TestPipeline
 
 try:
+  import tensorflow as tf
   from apache_beam.examples.inference import vertex_ai_image_classification
   from apache_beam.examples.inference import vertex_ai_llm_text_classification
 except ImportError as e:
@@ -45,6 +46,10 @@
 
 
 class VertexAIInference(unittest.TestCase):
+  @unittest.skipIf(
+    tf is None, 'Missing dependencies. '
+    'Test depends on tensorflow')
+  @pytest.mark.uses_tf
   @pytest.mark.uses_vertex_ai
   @pytest.mark.it_postcommit
   def test_vertex_ai_run_flower_image_classification(self):