From 2aa6361bc5735dd707ec8941baa9b971c8489292 Mon Sep 17 00:00:00 2001
From: Sergio Vidiella Pinto <sergiovidiella@google.com>
Date: Thu, 29 Feb 2024 13:17:05 +0100
Subject: [PATCH] fix: avoid init vertexai during unittests

---
 libs/vertexai/.gitignore                            |  1 +
 libs/vertexai/tests/integration_tests/conftest.py   |  2 +-
 .../tests/integration_tests/test_embeddings.py      |  6 +++---
 libs/vertexai/tests/unit_tests/test_embeddings.py   | 13 +++++++++++--
 4 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/libs/vertexai/.gitignore b/libs/vertexai/.gitignore
index bee8a64b..7dcfd0e2 100644
--- a/libs/vertexai/.gitignore
+++ b/libs/vertexai/.gitignore
@@ -1 +1,2 @@
 __pycache__
+.mypy_cache_test
diff --git a/libs/vertexai/tests/integration_tests/conftest.py b/libs/vertexai/tests/integration_tests/conftest.py
index feba4c07..b0903366 100644
--- a/libs/vertexai/tests/integration_tests/conftest.py
+++ b/libs/vertexai/tests/integration_tests/conftest.py
@@ -2,7 +2,7 @@
 
 import pytest
 from _pytest.tmpdir import TempPathFactory
-from vertexai.vision_models import Image
+from vertexai.vision_models import Image  # type: ignore
 
 
 @pytest.fixture
diff --git a/libs/vertexai/tests/integration_tests/test_embeddings.py b/libs/vertexai/tests/integration_tests/test_embeddings.py
index 210da903..56fb03cf 100644
--- a/libs/vertexai/tests/integration_tests/test_embeddings.py
+++ b/libs/vertexai/tests/integration_tests/test_embeddings.py
@@ -4,12 +4,12 @@
 `gcloud auth login` first).
 """
 import pytest
-from vertexai.language_models import TextEmbeddingModel
-from vertexai.vision_models import MultiModalEmbeddingModel
+from vertexai.language_models import TextEmbeddingModel  # type: ignore
+from vertexai.vision_models import MultiModalEmbeddingModel  # type: ignore
 
 from langchain_google_vertexai.embeddings import (
-    VertexAIEmbeddings,
     GoogleEmbeddingModelType,
+    VertexAIEmbeddings,
 )
 
 
diff --git a/libs/vertexai/tests/unit_tests/test_embeddings.py b/libs/vertexai/tests/unit_tests/test_embeddings.py
index 664e103a..8e1d86f4 100644
--- a/libs/vertexai/tests/unit_tests/test_embeddings.py
+++ b/libs/vertexai/tests/unit_tests/test_embeddings.py
@@ -1,7 +1,8 @@
-from typing import Any
+from typing import Any, Dict
 from unittest.mock import MagicMock
 
 import pytest
+from pydantic.v1 import root_validator
 
 from langchain_google_vertexai import VertexAIEmbeddings
 from langchain_google_vertexai.embeddings import GoogleEmbeddingModelType
@@ -39,4 +40,12 @@ class MockVertexAIEmbeddings(VertexAIEmbeddings):
 
     def __init__(self, model_name, **kwargs: Any) -> None:
         super().__init__(model_name, **kwargs)
-        self.client = MagicMock()
+
+    @classmethod
+    def _init_vertexai(cls, values: Dict) -> None:
+        pass
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        values["client"] = MagicMock()
+        return values