fix unit tests

Author: laurenyu

tests/unit/test_airflow.py

@@ -622,15 +622,12 @@ def test_framework_tuning_config(ecr_prefix, sagemaker_session):
         ]
     ),
 )
-@patch(
-    "sagemaker.fw_utils.get_ecr_image_uri_prefix",
-    return_value="520713654638.dkr.ecr.us-west-2.amazonaws.com",
-)
-@patch(
-    "sagemaker.amazon.amazon_estimator.get_ecr_image_uri_prefix",
-    return_value="174872318107.dkr.ecr.us-west-2.amazonaws.com",
-)
-def test_multi_estimator_tuning_config(algo_ecr_prefix, fw_ecr_prefix, sagemaker_session):
+@patch("sagemaker.utils._botocore_resolver")
+def test_multi_estimator_tuning_config(botocore_resolver, sagemaker_session):
+    botocore_resolver.return_value.construct_endpoint.return_value = {
+        "hostname": "ecr.us-west-2.amazonaws.com"
+    }
+
     estimator_dict = {}
     hyperparameter_ranges_dict = {}
     objective_metric_name_dict = {}

tests/unit/test_estimator.py

@@ -19,18 +19,17 @@
 from time import sleep
 
 import pytest
+from botocore.exceptions import ClientError
 from mock import ANY, MagicMock, Mock, patch
 
-from sagemaker import TrainingInput, utils, vpc_utils
-from sagemaker.amazon.amazon_estimator import registry
+import sagemaker.local
+from sagemaker import image_uris, TrainingInput, utils, vpc_utils
 from sagemaker.algorithm import AlgorithmEstimator
 from sagemaker.estimator import Estimator, EstimatorBase, Framework, _TrainingJob
 from sagemaker.model import FrameworkModel
 from sagemaker.predictor import Predictor
 from sagemaker.session import ShuffleConfig
 from sagemaker.transformer import Transformer
-from botocore.exceptions import ClientError
-import sagemaker.local
 
 MODEL_DATA = "s3://bucket/model.tar.gz"
 MODEL_IMAGE = "mi"
@@ -2373,10 +2372,9 @@ def test_prepare_for_training_with_name_based_on_algorithm(sagemaker_session):
     ),
 )
 def test_encryption_flag_in_non_vpc_mode_invalid(sagemaker_session):
-    image_uri = registry("us-west-2") + "/factorization-machines:1"
     with pytest.raises(ClientError) as error:
         estimator = Estimator(
-            image_uri=image_uri,
+            image_uri="some-image",
             role="SageMakerRole",
             instance_count=1,
             instance_type="ml.c4.xlarge",

tests/unit/test_fm.py

@@ -15,11 +15,12 @@
 import pytest
 from mock import Mock, patch
 
+from sagemaker import image_uris
 from sagemaker.amazon.factorization_machines import (
     FactorizationMachines,
     FactorizationMachinesPredictor,
 )
-from sagemaker.amazon.amazon_estimator import registry, RecordSet
+from sagemaker.amazon.amazon_estimator import RecordSet
 
 ROLE = "myrole"
 INSTANCE_COUNT = 1
@@ -146,7 +147,7 @@ def test_all_hyperparameters(sagemaker_session):
 
 def test_image(sagemaker_session):
     fm = FactorizationMachines(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
-    assert fm.train_image() == registry(REGION) + "/factorization-machines:1"
+    assert image_uris.retrieve("factorization-machines", REGION) == fm.train_image()
 
 
 @pytest.mark.parametrize(
@@ -313,9 +314,7 @@ def test_model_image(sagemaker_session):
     fm.fit(data, MINI_BATCH_SIZE)
 
     model = fm.create_model()
-    assert (
-        model.image_uri == registry(REGION, "factorization-machines") + "/factorization-machines:1"
-    )
+    assert image_uris.retrieve("factorization-machines", REGION) == model.image_uri
 
 
 def test_predictor_type(sagemaker_session):

tests/unit/test_ipinsights.py

@@ -15,8 +15,9 @@
 import pytest
 from mock import Mock, patch
 
+from sagemaker import image_uris
 from sagemaker.amazon.ipinsights import IPInsights, IPInsightsPredictor
-from sagemaker.amazon.amazon_estimator import registry, RecordSet
+from sagemaker.amazon.amazon_estimator import RecordSet
 
 # Mocked training config
 ROLE = "myrole"
@@ -119,7 +120,7 @@ def test_all_hyperparameters(sagemaker_session):
 
 def test_image(sagemaker_session):
     ipinsights = IPInsights(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
-    assert ipinsights.train_image() == registry(REGION, "ipinsights") + "/ipinsights:1"
+    assert image_uris.retrieve("ipinsights", REGION) == ipinsights.train_image()
 
 
 @pytest.mark.parametrize(
@@ -288,7 +289,7 @@ def test_model_image(sagemaker_session):
     ipinsights.fit(data, MINI_BATCH_SIZE)
 
     model = ipinsights.create_model()
-    assert model.image_uri == registry(REGION, "ipinsights") + "/ipinsights:1"
+    assert image_uris.retrieve("ipinsights", REGION) == model.image_uri
 
 
 def test_predictor_type(sagemaker_session):

tests/unit/test_kmeans.py

@@ -15,8 +15,9 @@
 import pytest
 from mock import Mock, patch
 
+from sagemaker import image_uris
 from sagemaker.amazon.kmeans import KMeans, KMeansPredictor
-from sagemaker.amazon.amazon_estimator import registry, RecordSet
+from sagemaker.amazon.amazon_estimator import RecordSet
 
 ROLE = "myrole"
 INSTANCE_COUNT = 1
@@ -111,7 +112,7 @@ def test_all_hyperparameters(sagemaker_session):
 
 def test_image(sagemaker_session):
     kmeans = KMeans(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
-    assert kmeans.train_image() == registry(REGION, "kmeans") + "/kmeans:1"
+    assert image_uris.retrieve("kmeans", REGION) == kmeans.train_image()
 
 
 @pytest.mark.parametrize("required_hyper_parameters, value", [("k", "string")])
@@ -255,7 +256,7 @@ def test_model_image(sagemaker_session):
     kmeans.fit(data, MINI_BATCH_SIZE)
 
     model = kmeans.create_model()
-    assert model.image_uri == registry(REGION, "kmeans") + "/kmeans:1"
+    assert image_uris.retrieve("kmeans", REGION) == model.image_uri
 
 
 def test_predictor_type(sagemaker_session):

tests/unit/test_knn.py

@@ -15,8 +15,9 @@
 import pytest
 from mock import Mock, patch
 
+from sagemaker import image_uris
 from sagemaker.amazon.knn import KNN, KNNPredictor
-from sagemaker.amazon.amazon_estimator import registry, RecordSet
+from sagemaker.amazon.amazon_estimator import RecordSet
 
 ROLE = "myrole"
 INSTANCE_COUNT = 1
@@ -145,7 +146,7 @@ def test_all_hyperparameters_classifier(sagemaker_session):
 
 def test_image(sagemaker_session):
     knn = KNN(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
-    assert knn.train_image() == registry(REGION, "knn") + "/knn:1"
+    assert image_uris.retrieve("knn", REGION) == knn.train_image()
 
 
 @pytest.mark.parametrize(
@@ -279,7 +280,7 @@ def test_model_image(sagemaker_session):
     knn.fit(data, MINI_BATCH_SIZE)
 
     model = knn.create_model()
-    assert model.image_uri == registry(REGION, "knn") + "/knn:1"
+    assert image_uris.retrieve("knn", REGION) == model.image_uri
 
 
 def test_predictor_type(sagemaker_session):

tests/unit/test_lda.py

@@ -15,8 +15,9 @@
 import pytest
 from mock import Mock, patch
 
+from sagemaker import image_uris
 from sagemaker.amazon.lda import LDA, LDAPredictor
-from sagemaker.amazon.amazon_estimator import registry, RecordSet
+from sagemaker.amazon.amazon_estimator import RecordSet
 
 ROLE = "myrole"
 INSTANCE_COUNT = 1
@@ -95,7 +96,7 @@ def test_all_hyperparameters(sagemaker_session):
 
 def test_image(sagemaker_session):
     lda = LDA(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
-    assert lda.train_image() == registry(REGION, "lda") + "/lda:1"
+    assert image_uris.retrieve("lda", REGION) == lda.train_image()
 
 
 @pytest.mark.parametrize("required_hyper_parameters, value", [("num_topics", "string")])
@@ -215,7 +216,7 @@ def test_model_image(sagemaker_session):
     lda.fit(data, MINI_BATCH_SZIE)
 
     model = lda.create_model()
-    assert model.image_uri == registry(REGION, "lda") + "/lda:1"
+    assert image_uris.retrieve("lda", REGION) == model.image_uri
 
 
 def test_predictor_type(sagemaker_session):

tests/unit/test_linear_learner.py

@@ -15,8 +15,9 @@
 import pytest
 from mock import Mock, patch
 
+from sagemaker import image_uris
 from sagemaker.amazon.linear_learner import LinearLearner, LinearLearnerPredictor
-from sagemaker.amazon.amazon_estimator import registry, RecordSet
+from sagemaker.amazon.amazon_estimator import RecordSet
 
 ROLE = "myrole"
 INSTANCE_COUNT = 1
@@ -178,7 +179,7 @@ def test_all_hyperparameters(sagemaker_session):
 
 def test_image(sagemaker_session):
     lr = LinearLearner(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
-    assert lr.train_image() == registry(REGION, "linear-learner") + "/linear-learner:1"
+    assert image_uris.retrieve("linear-learner", REGION) == lr.train_image()
 
 
 @pytest.mark.parametrize("required_hyper_parameters, value", [("predictor_type", 0)])
@@ -412,7 +413,7 @@ def test_model_image(sagemaker_session):
     lr.fit(data)
 
     model = lr.create_model()
-    assert model.image_uri == registry(REGION, "linear-learner") + "/linear-learner:1"
+    assert image_uris.retrieve("linear-learner", REGION) == model.image_uri
 
 
 def test_predictor_type(sagemaker_session):

tests/unit/test_ntm.py

@@ -15,8 +15,9 @@
 import pytest
 from mock import Mock, patch
 
+from sagemaker import image_uris
 from sagemaker.amazon.ntm import NTM, NTMPredictor
-from sagemaker.amazon.amazon_estimator import registry, RecordSet
+from sagemaker.amazon.amazon_estimator import RecordSet
 
 ROLE = "myrole"
 INSTANCE_COUNT = 1
@@ -114,7 +115,7 @@ def test_all_hyperparameters(sagemaker_session):
 
 def test_image(sagemaker_session):
     ntm = NTM(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
-    assert ntm.train_image() == registry(REGION, "ntm") + "/ntm:1"
+    assert image_uris.retrieve("ntm", REGION) == ntm.train_image()
 
 
 @pytest.mark.parametrize("required_hyper_parameters, value", [("num_topics", "string")])
@@ -278,7 +279,7 @@ def test_model_image(sagemaker_session):
     ntm.fit(data, MINI_BATCH_SIZE)
 
     model = ntm.create_model()
-    assert model.image_uri == registry(REGION, "ntm") + "/ntm:1"
+    assert image_uris.retrieve("ntm", REGION) == model.image_uri
 
 
 def test_predictor_type(sagemaker_session):

tests/unit/test_object2vec.py

@@ -15,9 +15,10 @@
 import pytest
 from mock import Mock, patch
 
+from sagemaker import image_uris
 from sagemaker.amazon.object2vec import Object2Vec
 from sagemaker.predictor import Predictor
-from sagemaker.amazon.amazon_estimator import registry, RecordSet
+from sagemaker.amazon.amazon_estimator import RecordSet
 
 ROLE = "myrole"
 INSTANCE_COUNT = 1
@@ -143,7 +144,7 @@ def test_all_hyperparameters(sagemaker_session):
 
 def test_image(sagemaker_session):
     object2vec = Object2Vec(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
-    assert object2vec.train_image() == registry(REGION, "object2vec") + "/object2vec:1"
+    assert image_uris.retrieve("object2vec", REGION) == object2vec.train_image()
 
 
 @pytest.mark.parametrize("required_hyper_parameters, value", [("epochs", "string")])
@@ -308,7 +309,7 @@ def test_model_image(sagemaker_session):
     object2vec.fit(data, MINI_BATCH_SIZE)
 
     model = object2vec.create_model()
-    assert model.image_uri == registry(REGION, "object2vec") + "/object2vec:1"
+    assert image_uris.retrieve("object2vec", REGION) == model.image_uri
 
 
 def test_predictor_type(sagemaker_session):

tests/unit/test_pca.py

@@ -15,8 +15,9 @@
 import pytest
 from mock import Mock, patch
 
+from sagemaker import image_uris
 from sagemaker.amazon.pca import PCA, PCAPredictor
-from sagemaker.amazon.amazon_estimator import registry, RecordSet
+from sagemaker.amazon.amazon_estimator import RecordSet
 
 ROLE = "myrole"
 INSTANCE_COUNT = 1
@@ -100,7 +101,7 @@ def test_all_hyperparameters(sagemaker_session):
 
 def test_image(sagemaker_session):
     pca = PCA(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
-    assert pca.train_image() == registry(REGION, "pca") + "/pca:1"
+    assert image_uris.retrieve("pca", REGION) == pca.train_image()
 
 
 @pytest.mark.parametrize("required_hyper_parameters, value", [("num_components", "string")])
@@ -231,7 +232,7 @@ def test_model_image(sagemaker_session):
     pca.fit(data, MINI_BATCH_SIZE)
 
     model = pca.create_model()
-    assert model.image_uri == registry(REGION, "pca") + "/pca:1"
+    assert image_uris.retrieve("pca", REGION) == model.image_uri
 
 
 def test_predictor_type(sagemaker_session):

tests/unit/test_randomcutforest.py

@@ -15,8 +15,9 @@
 import pytest
 from mock import Mock, patch
 
+from sagemaker import image_uris
 from sagemaker.amazon.randomcutforest import RandomCutForest, RandomCutForestPredictor
-from sagemaker.amazon.amazon_estimator import registry, RecordSet
+from sagemaker.amazon.amazon_estimator import RecordSet
 
 ROLE = "myrole"
 INSTANCE_COUNT = 1
@@ -106,9 +107,7 @@ def test_all_hyperparameters(sagemaker_session):
 
 def test_image(sagemaker_session):
     randomcutforest = RandomCutForest(sagemaker_session=sagemaker_session, **ALL_REQ_ARGS)
-    assert (
-        randomcutforest.train_image() == registry(REGION, "randomcutforest") + "/randomcutforest:1"
-    )
+    assert image_uris.retrieve("randomcutforest", REGION) == randomcutforest.train_image()
 
 
 @pytest.mark.parametrize("iterable_hyper_parameters, value", [("eval_metrics", 0)])
@@ -232,7 +231,7 @@ def test_model_image(sagemaker_session):
     randomcutforest.fit(data, MINI_BATCH_SIZE)
 
     model = randomcutforest.create_model()
-    assert model.image_uri == registry(REGION, "randomcutforest") + "/randomcutforest:1"
+    assert image_uris.retrieve("randomcutforest", REGION) == model.image_uri
 
 
 def test_predictor_type(sagemaker_session):