From d46d7ecf5c730697bbda9c81ddb8a226c7f7af52 Mon Sep 17 00:00:00 2001
From: Jarek Potiuk <jarek@potiuk.com>
Date: Fri, 29 Mar 2024 14:03:32 +0100
Subject: [PATCH] Remove google-cloud-aiplatform when tests without pydantic
 are run (#38617)

The #37818 bumped requirement for google-cloud-aiplatform to a
version that requires Pydantic so we need to remove it when
we are running "no-pydantic" tests.

The tests that import it should be skipped
when the dependency is removed.

(cherry picked from commit 69df34e52d01b7a75065ab2827152348d885ce0a)
---
 Dockerfile.ci                                                 | 2 +-
 pyproject.toml                                                | 2 ++
 scripts/docker/entrypoint_ci.sh                               | 2 +-
 .../google/cloud/hooks/vertex_ai/test_batch_prediction_job.py | 4 ++++
 .../cloud/hooks/vertex_ai/test_hyperparameter_tuning_job.py   | 4 ++++
 5 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/Dockerfile.ci b/Dockerfile.ci
index 694148ab50704..fcd2a40068d86 100644
--- a/Dockerfile.ci
+++ b/Dockerfile.ci
@@ -1021,7 +1021,7 @@ function check_pydantic() {
         echo
         # shellcheck disable=SC2086
         ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} pydantic aws-sam-translator openai \
-           pyiceberg qdrant-client cfn-lint weaviate-client
+           pyiceberg qdrant-client cfn-lint weaviate-client google-cloud-aiplatform
         pip check
     elif [[ ${PYDANTIC=} == "v1" ]]; then
         echo
diff --git a/pyproject.toml b/pyproject.toml
index 416f86060258e..737a65819794b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -355,6 +355,8 @@ combine-as-imports = true
 "airflow/security/utils.py" = ["E402"]
 "tests/providers/elasticsearch/log/elasticmock/__init__.py" = ["E402"]
 "tests/providers/elasticsearch/log/elasticmock/utilities/__init__.py" = ["E402"]
+"tests/providers/google/cloud/hooks/vertex_ai/test_batch_prediction_job.py" = ["E402"]
+"tests/providers/google/cloud/hooks/vertex_ai/test_hyperparameter_tuning_job.py" = ["E402"]
 "tests/providers/openai/hooks/test_openai.py" = ["E402"]
 "tests/providers/openai/operators/test_openai.py" = ["E402"]
 "tests/providers/qdrant/hooks/test_qdrant.py" = ["E402"]
diff --git a/scripts/docker/entrypoint_ci.sh b/scripts/docker/entrypoint_ci.sh
index 8b2ec5c09fad8..a0738dc1d6b8c 100755
--- a/scripts/docker/entrypoint_ci.sh
+++ b/scripts/docker/entrypoint_ci.sh
@@ -243,7 +243,7 @@ function check_pydantic() {
         echo
         # shellcheck disable=SC2086
         ${PACKAGING_TOOL_CMD} uninstall ${EXTRA_UNINSTALL_FLAGS} pydantic aws-sam-translator openai \
-           pyiceberg qdrant-client cfn-lint weaviate-client
+           pyiceberg qdrant-client cfn-lint weaviate-client google-cloud-aiplatform
         pip check
     elif [[ ${PYDANTIC=} == "v1" ]]; then
         echo
diff --git a/tests/providers/google/cloud/hooks/vertex_ai/test_batch_prediction_job.py b/tests/providers/google/cloud/hooks/vertex_ai/test_batch_prediction_job.py
index 33b83bc352930..1b2e0e1fbb9eb 100644
--- a/tests/providers/google/cloud/hooks/vertex_ai/test_batch_prediction_job.py
+++ b/tests/providers/google/cloud/hooks/vertex_ai/test_batch_prediction_job.py
@@ -22,6 +22,10 @@
 
 import pytest
 from google.api_core.gapic_v1.method import DEFAULT
+
+# For no Pydantic environment, we need to skip the tests
+pytest.importorskip("google.cloud.aiplatform_v1")
+
 from google.cloud.aiplatform_v1 import JobState
 
 from airflow.exceptions import AirflowException
diff --git a/tests/providers/google/cloud/hooks/vertex_ai/test_hyperparameter_tuning_job.py b/tests/providers/google/cloud/hooks/vertex_ai/test_hyperparameter_tuning_job.py
index ce71b5f9078b0..d2609b7f2faf6 100644
--- a/tests/providers/google/cloud/hooks/vertex_ai/test_hyperparameter_tuning_job.py
+++ b/tests/providers/google/cloud/hooks/vertex_ai/test_hyperparameter_tuning_job.py
@@ -22,6 +22,10 @@
 import pytest
 from google.api_core.gapic_v1.method import DEFAULT
 from google.cloud import aiplatform
+
+# For no Pydantic environment, we need to skip the tests
+pytest.importorskip("google.cloud.aiplatform_v1")
+
 from google.cloud.aiplatform_v1 import JobState
 
 from airflow.exceptions import AirflowException