docs: change init sample to use vertexai

sasha-gitg · copybara-github · commit 829e0f6fd286 · 2024-07-26T09:43:17.000-07:00
PiperOrigin-RevId: 656432815
diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py
@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock
+from unittest.mock import patch
 
 from google.cloud import aiplatform
+import vertexai
 from vertexai.resources import preview as preview_resources
 import pytest
 
@@ -25,6 +27,12 @@ def mock_sdk_init():
         yield mock
 
 
+@pytest.fixture
+def mock_vertexai_init():
+    with patch.object(vertexai, "init") as mock:
+        yield mock
+
+
 """
 ----------------------------------------------------------------------------
 Dataset Fixtures
diff --git a/samples/model-builder/init_sample.py b/samples/model-builder/init_sample.py
@@ -28,9 +28,9 @@ def init_sample(
     service_account: Optional[str] = None,
 ):
 
-    from google.cloud import aiplatform
+    import vertexai
 
-    aiplatform.init(
+    vertexai.init(
         project=project,
         location=location,
         experiment=experiment,
diff --git a/samples/model-builder/init_sample_test.py b/samples/model-builder/init_sample_test.py
@@ -17,7 +17,7 @@
 import test_constants as constants
 
 
-def test_init_sample(mock_sdk_init):
+def test_init_sample(mock_vertexai_init):
 
     init_sample.init_sample(
         project=constants.PROJECT,
@@ -29,7 +29,7 @@ def test_init_sample(mock_sdk_init):
         service_account=constants.SERVICE_ACCOUNT,
     )
 
-    mock_sdk_init.assert_called_once_with(
+    mock_vertexai_init.assert_called_once_with(
         project=constants.PROJECT,
         location=constants.LOCATION_EUROPE,
         experiment=constants.EXPERIMENT_NAME,
