Fixed tests

ivanmkc · ivanmkc · commit 348dae93abe2 · 2021-04-30T14:26:10.000-04:00
diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py
@@ -218,23 +218,14 @@ def mock_model():
     mock = MagicMock(aiplatform.models.Model)
     yield mock
 
+
 @pytest.fixture
-<<<<<<< HEAD
-def mock_get_model(mock_model):
-    with patch.object(aiplatform, "Model") as mock:
-        mock.return_value = mock_model
-=======
 def mock_init_model(mock_model):
-<<<<<<< HEAD
-    with patch.object(aiplatform.models.Model, "__init__") as mock:
-        mock.return_value = None
->>>>>>> 190388a (Fixed tests)
-=======
     with patch.object(aiplatform, "Model") as mock:
         mock.return_value = mock_model
->>>>>>> 1f2700d (Fixed tests)
         yield mock
 
+
 @pytest.fixture
 def mock_batch_predict_model(mock_model):
     with patch.object(mock_model, "batch_predict") as mock:
diff --git a/samples/model-builder/create_batch_prediction_job_sample_test.py b/samples/model-builder/create_batch_prediction_job_sample_test.py
@@ -18,11 +18,7 @@
 
 
 def test_create_batch_prediction_job_sample(
-<<<<<<< HEAD
-    mock_sdk_init, mock_get_model, mock_batch_predict_model
-=======
     mock_sdk_init, mock_model, mock_init_model, mock_batch_predict_model
->>>>>>> 1f2700d (Fixed tests)
 ):
 
     create_batch_prediction_job_sample.create_batch_prediction_job_sample(
@@ -37,7 +33,7 @@ def test_create_batch_prediction_job_sample(
     mock_sdk_init.assert_called_once_with(
         project=constants.PROJECT, location=constants.LOCATION
     )
-    mock_get_model.assert_called_once_with(constants.MODEL_NAME)
+    mock_init_model.assert_called_once_with(constants.MODEL_NAME)
     mock_batch_predict_model.assert_called_once_with(
         job_display_name=constants.DISPLAY_NAME,
         gcs_source=constants.GCS_SOURCES,
diff --git a/samples/model-builder/deploy_model_with_automatic_resources_test.py b/samples/model-builder/deploy_model_with_automatic_resources_test.py
@@ -31,7 +31,7 @@ def test_deploy_model_with_automatic_resources_sample(
         traffic_split=constants.TRAFFIC_SPLIT,
         min_replica_count=constants.MIN_REPLICA_COUNT,
         max_replica_count=constants.MAX_REPLICA_COUNT,
-        metadata=constants.METADATA,
+        metadata=constants.ENDPOINT_DEPLOY_METADATA,
     )
 
     mock_sdk_init.assert_called_once_with(
@@ -47,6 +47,6 @@ def test_deploy_model_with_automatic_resources_sample(
         traffic_split=constants.TRAFFIC_SPLIT,
         min_replica_count=constants.MIN_REPLICA_COUNT,
         max_replica_count=constants.MAX_REPLICA_COUNT,
-        metadata=constants.METADATA,
+        metadata=constants.ENDPOINT_DEPLOY_METADATA,
         sync=True,
     )
diff --git a/samples/model-builder/deploy_model_with_dedicated_resources_test.py b/samples/model-builder/deploy_model_with_dedicated_resources_test.py
@@ -36,7 +36,7 @@ def test_deploy_model_with_dedicated_resources_sample(
         accelerator_count=constants.ACCELERATOR_COUNT,
         explanation_metadata=constants.EXPLANATION_METADATA,
         explanation_parameters=constants.EXPLANATION_PARAMETERS,
-        metadata=constants.METADATA,
+        metadata=constants.ENDPOINT_DEPLOY_METADATA,
     )
 
     mock_sdk_init.assert_called_once_with(
@@ -57,6 +57,6 @@ def test_deploy_model_with_dedicated_resources_sample(
         accelerator_count=constants.ACCELERATOR_COUNT,
         explanation_metadata=constants.EXPLANATION_METADATA,
         explanation_parameters=constants.EXPLANATION_PARAMETERS,
-        metadata=constants.METADATA,
+        metadata=constants.ENDPOINT_DEPLOY_METADATA,
         sync=True,
     )
diff --git a/samples/model-builder/test_constants.py b/samples/model-builder/test_constants.py
@@ -148,3 +148,4 @@
 TRAFFIC_SPLIT = {"a": 99, "b": 1}
 MIN_REPLICA_COUNT = 1
 MAX_REPLICA_COUNT = 1
+ENDPOINT_DEPLOY_METADATA = ()
