From 8cdc7416e4c2517fc3f6f11611a27b4d8b170e4b Mon Sep 17 00:00:00 2001
From: secastillo <sarhcas@amazon.com>
Date: Fri, 25 Aug 2023 19:49:29 -0400
Subject: [PATCH] change: Add Data Wrangler TLV and version 3.x images

---
 .../image_uri_config/data-wrangler.json       | 29 +++++++++++++++
 .../image_uris/test_data_wrangler.py          | 35 ++++++++++++++-----
 .../sagemaker/wrangler/test_processing.py     |  2 +-
 3 files changed, 57 insertions(+), 9 deletions(-)

diff --git a/src/sagemaker/image_uri_config/data-wrangler.json b/src/sagemaker/image_uri_config/data-wrangler.json
index b25eb8e1bf..8f249b0056 100644
--- a/src/sagemaker/image_uri_config/data-wrangler.json
+++ b/src/sagemaker/image_uri_config/data-wrangler.json
@@ -56,6 +56,35 @@
           "cn-northwest-1": "249157047649"
         },
         "repository": "sagemaker-data-wrangler-container"
+      },
+      "3.x": {
+        "registries": {
+          "af-south-1": "143210264188",
+          "ap-east-1": "707077482487",
+          "ap-northeast-1": "649008135260",
+          "ap-northeast-2": "131546521161",
+          "ap-northeast-3": "913387583493",
+          "ap-south-1": "089933028263",
+          "ap-southeast-1": "119527597002",
+          "ap-southeast-2": "422173101802",
+          "ca-central-1": "557239378090",
+          "eu-central-1": "024640144536",
+          "eu-north-1": "054986407534",
+          "eu-south-1": "488287956546",
+          "eu-west-1": "245179582081",
+          "eu-west-2": "894491911112",
+          "eu-west-3": "807237891255",
+          "il-central-1": "406833011540",
+          "me-south-1": "376037874950",
+          "sa-east-1": "424196993095",
+          "us-east-1": "663277389841",
+          "us-east-2": "415577184552",
+          "us-west-1": "926135532090",
+          "us-west-2": "174368400705",
+          "cn-north-1": "245909111842",
+          "cn-northwest-1": "249157047649"
+        },
+        "repository": "sagemaker-data-wrangler-container"
       }
     }
   }
diff --git a/tests/unit/sagemaker/image_uris/test_data_wrangler.py b/tests/unit/sagemaker/image_uris/test_data_wrangler.py
index 0252d2b71a..93852acedd 100644
--- a/tests/unit/sagemaker/image_uris/test_data_wrangler.py
+++ b/tests/unit/sagemaker/image_uris/test_data_wrangler.py
@@ -40,20 +40,39 @@
     "cn-north-1": "245909111842",
     "cn-northwest-1": "249157047649",
 }
-VERSIONS = ["1.x", "2.x"]
+
+# Accounts only supported in DW 3.x and beyond
+DATA_WRANGLER_3X_ACCOUNTS = {
+    "il-central-1": "406833011540",
+}
+
+VERSIONS = ["1.x", "2.x", "3.x"]
+
+
+def _test_ecr_uri(account, region, version):
+    actual_uri = image_uris.retrieve("data-wrangler", region=region, version=version)
+    expected_uri = expected_uris.algo_uri(
+        "sagemaker-data-wrangler-container",
+        account,
+        region,
+        version=version,
+    )
+    return expected_uri == actual_uri
 
 
 def test_data_wrangler_ecr_uri():
     for version in VERSIONS:
         for region in DATA_WRANGLER_ACCOUNTS.keys():
-            actual_uri = image_uris.retrieve("data-wrangler", region=region, version="1.x")
-            expected_uri = expected_uris.algo_uri(
-                "sagemaker-data-wrangler-container",
-                DATA_WRANGLER_ACCOUNTS[region],
-                region,
-                version="1.x",
+            assert _test_ecr_uri(
+                account=DATA_WRANGLER_ACCOUNTS[region], region=region, version=version
             )
-            assert expected_uri == actual_uri
+
+
+def test_data_wrangler_ecr_uri_3x():
+    for region in DATA_WRANGLER_3X_ACCOUNTS.keys():
+        assert _test_ecr_uri(
+            account=DATA_WRANGLER_3X_ACCOUNTS[region], region=region, version="3.x"
+        )
 
 
 def test_data_wrangler_ecr_uri_none():
diff --git a/tests/unit/sagemaker/wrangler/test_processing.py b/tests/unit/sagemaker/wrangler/test_processing.py
index 37e13aff6e..01e85543e6 100644
--- a/tests/unit/sagemaker/wrangler/test_processing.py
+++ b/tests/unit/sagemaker/wrangler/test_processing.py
@@ -23,7 +23,7 @@
 REGION = "us-west-2"
 DATA_WRANGLER_RECIPE_SOURCE = "s3://data_wrangler_flows/flow-26-18-43-16-0b48ac2e.flow"
 DATA_WRANGLER_CONTAINER_URI = (
-    "174368400705.dkr.ecr.us-west-2.amazonaws.com/sagemaker-data-wrangler-container:2.x"
+    "174368400705.dkr.ecr.us-west-2.amazonaws.com/sagemaker-data-wrangler-container:3.x"
 )
 MOCK_S3_URI = "s3://mock_data/mock.csv"