Skip to content

Commit 410ab2c

Browse files
knikureknikure
authored
change: Enhance unit-tests to automatically consume image URIs config registries from config JSONs (#4126)
1 parent 69deaae commit 410ab2c

File tree

4 files changed

+54
-101
lines changed

4 files changed

+54
-101
lines changed

‎tests/unit/sagemaker/image_uris/conftest.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
from __future__ import absolute_import
14+
15+
import os
16+
import json
17+
import pytest
18+
19+
20+
@pytest.fixture(scope="module")
21+
def config_dir():
22+
return "src/sagemaker/image_uri_config/"
23+
24+
25+
@pytest.fixture(scope="module")
26+
def load_config(config_dir, request):
27+
config_file_name = request.param
28+
config_file_path = os.path.join(config_dir, config_file_name)
29+
30+
with open(config_file_path, "r") as config_file:
31+
return json.load(config_file)
32+
33+
34+
@pytest.fixture(scope="module")
35+
def extract_versions_for_image_scope(load_config, request):
36+
scope_val = request.param
37+
return load_config[scope_val]["versions"]

‎tests/unit/sagemaker/image_uris/test_base_python.py‎

Lines changed: 3 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -15,39 +15,11 @@
1515
from sagemaker import image_uris
1616
from tests.unit.sagemaker.image_uris import expected_uris
1717

18-
REGISTRIES = {
19-
"us-east-2": "429704687514",
20-
"me-south-1": "117516905037",
21-
"us-west-2": "236514542706",
22-
"ca-central-1": "310906938811",
23-
"ap-east-1": "493642496378",
24-
"us-east-1": "081325390199",
25-
"ap-northeast-2": "806072073708",
26-
"eu-west-2": "712779665605",
27-
"ap-southeast-2": "52832661640",
28-
"cn-northwest-1": "390780980154",
29-
"eu-north-1": "243637512696",
30-
"cn-north-1": "390048526115",
31-
"ap-south-1": "394103062818",
32-
"eu-west-3": "615547856133",
33-
"ap-southeast-3": "276181064229",
34-
"af-south-1": "559312083959",
35-
"eu-west-1": "470317259841",
36-
"eu-central-1": "936697816551",
37-
"sa-east-1": "782484402741",
38-
"ap-northeast-3": "792733760839",
39-
"eu-south-1": "592751261982",
40-
"ap-northeast-1": "102112518831",
41-
"us-west-1": "742091327244",
42-
"ap-southeast-1": "492261229750",
43-
"me-central-1": "103105715889",
44-
"us-gov-east-1": "107072934176",
45-
"us-gov-west-1": "107173498710",
46-
}
47-
4818

19+
@pytest.mark.parametrize("load_config", ["sagemaker-base-python.json"], indirect=True)
4920
@pytest.mark.parametrize("py_version", ["310", "38"])
50-
def test_get_base_python_image_uri(py_version):
21+
def test_get_base_python_image_uri(py_version, load_config):
22+
REGISTRIES = load_config["versions"]["1.0"]["registries"]
5123
for region in REGISTRIES.keys():
5224
uri = image_uris.get_base_python_image_uri(
5325
region=region,

‎tests/unit/sagemaker/image_uris/test_data_wrangler.py‎

Lines changed: 10 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -12,42 +12,10 @@
1212
# language governing permissions and limitations under the License.
1313
from __future__ import absolute_import
1414

15+
import pytest
1516
from sagemaker import image_uris
1617
from tests.unit.sagemaker.image_uris import expected_uris
1718

18-
DATA_WRANGLER_ACCOUNTS = {
19-
"af-south-1": "143210264188",
20-
"ap-east-1": "707077482487",
21-
"ap-northeast-1": "649008135260",
22-
"ap-northeast-2": "131546521161",
23-
"ap-northeast-3": "913387583493",
24-
"ap-south-1": "089933028263",
25-
"ap-southeast-1": "119527597002",
26-
"ap-southeast-2": "422173101802",
27-
"ca-central-1": "557239378090",
28-
"eu-central-1": "024640144536",
29-
"eu-north-1": "054986407534",
30-
"eu-south-1": "488287956546",
31-
"eu-west-1": "245179582081",
32-
"eu-west-2": "894491911112",
33-
"eu-west-3": "807237891255",
34-
"me-south-1": "376037874950",
35-
"sa-east-1": "424196993095",
36-
"us-east-1": "663277389841",
37-
"us-east-2": "415577184552",
38-
"us-west-1": "926135532090",
39-
"us-west-2": "174368400705",
40-
"cn-north-1": "245909111842",
41-
"cn-northwest-1": "249157047649",
42-
}
43-
44-
# Accounts only supported in DW 3.x and beyond
45-
DATA_WRANGLER_3X_ACCOUNTS = {
46-
"il-central-1": "406833011540",
47-
}
48-
49-
VERSIONS = ["1.x", "2.x", "3.x"]
50-
5119

5220
def _test_ecr_uri(account, region, version):
5321
actual_uri = image_uris.retrieve("data-wrangler", region=region, version=version)
@@ -60,23 +28,23 @@ def _test_ecr_uri(account, region, version):
6028
return expected_uri == actual_uri
6129

6230

63-
def test_data_wrangler_ecr_uri():
31+
@pytest.mark.parametrize("load_config", ["data-wrangler.json"], indirect=True)
32+
@pytest.mark.parametrize("extract_versions_for_image_scope", ["processing"], indirect=True)
33+
def test_data_wrangler_ecr_uri(load_config, extract_versions_for_image_scope):
34+
VERSIONS = extract_versions_for_image_scope
6435
for version in VERSIONS:
36+
DATA_WRANGLER_ACCOUNTS = load_config["processing"]["versions"][version]["registries"]
6537
for region in DATA_WRANGLER_ACCOUNTS.keys():
6638
assert _test_ecr_uri(
6739
account=DATA_WRANGLER_ACCOUNTS[region], region=region, version=version
6840
)
6941

7042

71-
def test_data_wrangler_ecr_uri_3x():
72-
for region in DATA_WRANGLER_3X_ACCOUNTS.keys():
73-
assert _test_ecr_uri(
74-
account=DATA_WRANGLER_3X_ACCOUNTS[region], region=region, version="3.x"
75-
)
76-
77-
78-
def test_data_wrangler_ecr_uri_none():
43+
@pytest.mark.parametrize("load_config", ["data-wrangler.json"], indirect=True)
44+
def test_data_wrangler_ecr_uri_none(load_config):
7945
region = "us-west-2"
46+
VERSIONS = ["1.x", "2.x", "3.x"]
47+
DATA_WRANGLER_ACCOUNTS = load_config["processing"]["versions"]["1.x"]["registries"]
8048
actual_uri = image_uris.retrieve("data-wrangler", region=region)
8149
expected_uri = expected_uris.algo_uri(
8250
"sagemaker-data-wrangler-container",

‎tests/unit/sagemaker/image_uris/test_debugger.py‎

Lines changed: 4 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -12,38 +12,14 @@
1212
# language governing permissions and limitations under the License.
1313
from __future__ import absolute_import
1414

15+
import pytest
1516
from sagemaker import image_uris
1617
from tests.unit.sagemaker.image_uris import expected_uris
1718

18-
ACCOUNTS = {
19-
"af-south-1": "314341159256",
20-
"ap-east-1": "199566480951",
21-
"ap-northeast-1": "430734990657",
22-
"ap-northeast-2": "578805364391",
23-
"ap-northeast-3": "479947661362",
24-
"ap-south-1": "904829902805",
25-
"ap-southeast-1": "972752614525",
26-
"ap-southeast-2": "184798709955",
27-
"ca-central-1": "519511493484",
28-
"cn-north-1": "618459771430",
29-
"cn-northwest-1": "658757709296",
30-
"eu-central-1": "482524230118",
31-
"eu-north-1": "314864569078",
32-
"eu-south-1": "563282790590",
33-
"eu-west-1": "929884845733",
34-
"eu-west-2": "250201462417",
35-
"eu-west-3": "447278800020",
36-
"me-south-1": "986000313247",
37-
"sa-east-1": "818342061345",
38-
"us-east-1": "503895931360",
39-
"us-east-2": "915447279597",
40-
"us-gov-west-1": "515509971035",
41-
"us-west-1": "685455198987",
42-
"us-west-2": "895741380848",
43-
}
4419

45-
46-
def test_debugger():
20+
@pytest.mark.parametrize("load_config", ["debugger.json"], indirect=True)
21+
def test_debugger(load_config):
22+
ACCOUNTS = load_config["versions"]["latest"]["registries"]
4723
for region in ACCOUNTS.keys():
4824
uri = image_uris.retrieve("debugger", region=region)
4925
expected = expected_uris.algo_uri(

0 commit comments

Comments
 (0)