From 32ef90b1669be35bd623034771df62decea31f92 Mon Sep 17 00:00:00 2001 From: ivanmkc Date: Fri, 16 Apr 2021 23:06:31 -0400 Subject: [PATCH] Fixed tests --- samples/model-builder/conftest.py | 8 +++++++- samples/model-builder/deploy_model_sample.py | 2 +- samples/model-builder/deploy_model_test.py | 4 +++- samples/model-builder/test_constants.py | 1 + samples/model-builder/upload_model_test.py | 2 ++ 5 files changed, 14 insertions(+), 3 deletions(-) diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py index 950302b2b4..0414d53c1a 100644 --- 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 index db3fa17e5a..1d330f4059 100644 --- 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 index 51bcb7ebf6..708901dabf 100644 --- 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 index c5131d4217..11b4964608 100644 --- 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 index 35b84fc850..30f3ee8cce 100644 --- 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,