From 2fecc4ecdd7da22a9d83f8fc26033bd35f3bfc72 Mon Sep 17 00:00:00 2001
From: rd-pong <rdpeng@amazon.com>
Date: Thu, 13 Apr 2023 09:58:56 -0700
Subject: [PATCH] Create shallow canary marker Add placeholder values in
 replacement_values.py to avoid KeyError

---
 test/canary/scripts/run_test.sh        |  2 ++
 test/e2e/conftest.py                   |  1 +
 test/e2e/replacement_values.py         | 13 +++++++++++++
 test/e2e/tests/test_endpoint.py        |  1 +
 test/e2e/tests/test_endpoint_config.py |  1 +
 test/e2e/tests/test_trainingjob.py     |  1 +
 6 files changed, 19 insertions(+)

diff --git a/test/canary/scripts/run_test.sh b/test/canary/scripts/run_test.sh
index d5dc2ccb..f1b36016 100755
--- a/test/canary/scripts/run_test.sh
+++ b/test/canary/scripts/run_test.sh
@@ -97,6 +97,8 @@ pushd $E2E_DIR
   if [[ $SERVICE_REGION =~ ^(eu-north-1|eu-west-3)$  ]]; then
     # If select_regions_1 true we run the notebook_instance test
     pytest_args+=(-m "canary or select_regions_1")
+  elif [[ $SERVICE_REGION =~ ^(eu-south-2|ap-southeast-3|me-central-1|eu-central-2)$  ]]; then
+    pytest_args+=(-m "shallow_canary")
   else
     pytest_args+=(-m "canary")
   fi
diff --git a/test/e2e/conftest.py b/test/e2e/conftest.py
index c96fa8a1..9ec910a2 100644
--- a/test/e2e/conftest.py
+++ b/test/e2e/conftest.py
@@ -30,6 +30,7 @@ def pytest_configure(config):
     )
     config.addinivalue_line("markers", "slow: mark test as slow to run")
     config.addinivalue_line("markers", "select_regions_1: mark test to only run if in select region")
+    config.addinivalue_line("markers", "shallow_canary: mark test to run in shallow canary tests")
 
 
 def pytest_collection_modifyitems(config, items):
diff --git a/test/e2e/replacement_values.py b/test/e2e/replacement_values.py
index 89daffd6..09c209f4 100644
--- a/test/e2e/replacement_values.py
+++ b/test/e2e/replacement_values.py
@@ -47,6 +47,7 @@
     "ap-southeast-3": "951798379941.dkr.ecr.ap-southeast-3.amazonaws.com",
     "me-central-1": "272398656194.dkr.ecr.me-central-1.amazonaws.com",
     "eu-central-2": "680994064768.dkr.ecr.eu-central-2.amazonaws.com",
+    "af-south-1": "510948584623.dkr.ecr.af-south-1.amazonaws.com",
 }
 
 # https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/image_uri_config/debugger.json
@@ -73,6 +74,11 @@
     "sa-east-1": "818342061345.dkr.ecr.sa-east-1.amazonaws.com",
     "eu-south-1": "563282790590.dkr.ecr.eu-south-1.amazonaws.com",
     "ap-northeast-3": "479947661362.dkr.ecr.ap-northeast-3.amazonaws.com",
+    "af-south-1": "314341159256.dkr.ecr.af-south-1.amazonaws.com",
+    "eu-south-2": "",
+    "ap-southeast-3": "",
+    "me-central-1": "",
+    "eu-central-2": "",
 }
 
 # https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-algo-docker-registry-paths.html
@@ -103,6 +109,7 @@
     "ap-southeast-3": "951798379941.dkr.ecr.ap-southeast-3.amazonaws.com",
     "me-central-1": "272398656194.dkr.ecr.me-central-1.amazonaws.com",
     "eu-central-2": "680994064768.dkr.ecr.eu-central-2.amazonaws.com",
+    "af-south-1": "455444449433.dkr.ecr.af-south-1.amazonaws.com",
 }
 
 
@@ -163,6 +170,9 @@
     "us-gov-west-1": "362178532790.dkr.ecr.us-gov-west-1.amazonaws.com",
     "ap-northeast-3": "990339680094.dkr.ecr.ap-northeast-3.amazonaws.com",
     "ap-southeast-3": "669540362728.dkr.ecr.ap-southeast-3.amazonaws.com",
+    "eu-south-2": "",
+    "me-central-1": "",
+    "eu-central-2": "",
 }
 
 # https://docs.aws.amazon.com/sagemaker/latest/dg/clarify-configure-processing-jobs.html#clarify-processing-job-configure-container
@@ -189,6 +199,9 @@
     "eu-south-1": "638885417683.dkr.ecr.eu-south-1.amazonaws.com",
     "ap-northeast-3": "912233562940.dkr.ecr.ap-northeast-3.amazonaws.com",
     "ap-southeast-3": "705930551576.dkr.ecr.ap-southeast-3.amazonaws.com",
+    "eu-south-2": "",
+    "me-central-1": "",
+    "eu-central-2": "",
 }
 
 ENDPOINT_INSTANCE_TYPES = {
diff --git a/test/e2e/tests/test_endpoint.py b/test/e2e/tests/test_endpoint.py
index 4b8e2712..b973e44f 100644
--- a/test/e2e/tests/test_endpoint.py
+++ b/test/e2e/tests/test_endpoint.py
@@ -218,6 +218,7 @@ def faulty_config(name_suffix, single_container_model):
 
 
 @service_marker
+@pytest.mark.shallow_canary
 @pytest.mark.canary
 class TestEndpoint:
     def create_endpoint_test(self, xgboost_endpoint):
diff --git a/test/e2e/tests/test_endpoint_config.py b/test/e2e/tests/test_endpoint_config.py
index 33a040fb..55407f77 100644
--- a/test/e2e/tests/test_endpoint_config.py
+++ b/test/e2e/tests/test_endpoint_config.py
@@ -74,6 +74,7 @@ def single_variant_config():
 
 
 @service_marker
+@pytest.mark.shallow_canary
 @pytest.mark.canary
 class TestEndpointConfig:
     def test_create_endpoint_config(self, single_variant_config):
diff --git a/test/e2e/tests/test_trainingjob.py b/test/e2e/tests/test_trainingjob.py
index bfa4d5f3..057d5eb9 100644
--- a/test/e2e/tests/test_trainingjob.py
+++ b/test/e2e/tests/test_trainingjob.py
@@ -85,6 +85,7 @@ def test_stopped(self, xgboost_training_job):
         training_job_desc = get_sagemaker_training_job(training_job_name)
         assert training_job_desc["TrainingJobStatus"] in cfg.LIST_JOB_STATUS_STOPPED
 
+    @pytest.mark.shallow_canary
     @pytest.mark.canary
     def test_completed(self, xgboost_training_job):
         (reference, resource) = xgboost_training_job