Fixed tests

googleapis · Apr 17, 2021 · 32ef90b · 32ef90b
1 parent 11acc60
commit 32ef90b
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 3 deletions.
diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py
@@ -160,7 +160,13 @@ def mock_run_automl_image_training_job():
 
 
 @pytest.fixture
-def mock_init_model():
+def mock_model():
+    mock = MagicMock(aiplatform.models.Model)
+    yield mock
+
+
+@pytest.fixture
+def mock_init_model(mock_model):
     with patch.object(aiplatform.models.Model, "__init__") as mock:
         mock.return_value = None
         yield mock

diff --git a/samples/model-builder/deploy_model_sample.py b/samples/model-builder/deploy_model_sample.py
@@ -39,7 +39,7 @@ def deploy_model_sample(
 
     aiplatform.init(project=project, location=location)
 
-    model = aiplatform.Model(model_resource_name)
+    model = aiplatform.Model(model_name=model_resource_name)
 
     model.deploy(
         endpoint=endpoint,

diff --git a/samples/model-builder/deploy_model_test.py b/samples/model-builder/deploy_model_test.py
@@ -17,11 +17,12 @@
 import test_constants as constants
 
 
-def test_deploy_model_sample(mock_sdk_init, mock_deploy_model):
+def test_deploy_model_sample(mock_sdk_init, mock_init_model, mock_deploy_model):
 
     deploy_model_sample.deploy_model_sample(
         project=constants.PROJECT,
         location=constants.LOCATION,
+        model_resource_name=constants.MODEL_RESOURCE_NAME,
         endpoint=constants.ENDPOINT_NAME,
         deployed_model_display_name=constants.DEPLOYED_MODEL_DISPLAY_NAME,
         traffic_percentage=constants.TRAFFIC_PERCENTAGE,
@@ -42,6 +43,7 @@ def test_deploy_model_sample(mock_sdk_init, mock_deploy_model):
     )
 
     mock_deploy_model.assert_called_once_with(
+        model_resource_name=constants.MODEL_RESOURCE_NAME,
         endpoint=constants.ENDPOINT_NAME,
         deployed_model_display_name=constants.DEPLOYED_MODEL_DISPLAY_NAME,
         traffic_percentage=constants.TRAFFIC_PERCENTAGE,

diff --git a/samples/model-builder/test_constants.py b/samples/model-builder/test_constants.py
@@ -54,6 +54,7 @@
 PREDICTION_TEXT_INSTANCE = "This is some text for testing NLP prediction output"
 
 # Model constants
+MODEL_RESOURCE_NAME = f"{PARENT}/models/1234"
 MODEL_ARTIFACT_URI = "gs://bucket3/output-dir/"
 SERVING_CONTAINER_IMAGE_URI = "http://gcr.io/test/test:latest"
 SERVING_CONTAINER_IMAGE = "gcr.io/test-serving/container:image"

diff --git a/samples/model-builder/upload_model_test.py b/samples/model-builder/upload_model_test.py
@@ -22,6 +22,7 @@ def test_upload_model_sample(mock_sdk_init, mock_upload_model):
     upload_model_sample.upload_model_sample(
         project=constants.PROJECT,
         location=constants.LOCATION,
+        display_name=constants.MODEL_NAME,
         artifact_uri=constants.MODEL_ARTIFACT_URI,
         serving_container_image_uri=constants.SERVING_CONTAINER_IMAGE_URI,
         serving_container_predict_route=constants.SERVING_CONTAINER_PREDICT_ROUTE,
@@ -44,6 +45,7 @@ def test_upload_model_sample(mock_sdk_init, mock_upload_model):
     )
 
     mock_upload_model.assert_called_once_with(
+        display_name=constants.MODEL_NAME,
         artifact_uri=constants.MODEL_ARTIFACT_URI,
         serving_container_image_uri=constants.SERVING_CONTAINER_IMAGE_URI,
         serving_container_predict_route=constants.SERVING_CONTAINER_PREDICT_ROUTE,