docs: change init sample to use vertexai

PiperOrigin-RevId: 656432815

Author: sasha-gitg

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

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,

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,