fix dict KeyError and and unit tests

Author: Chuyang Deng

src/sagemaker/cli/framework_upgrade.py

@@ -22,7 +22,7 @@
 IMAGE_URI_CONFIG_DIR = os.path.join("..", "image_uri_config")
 
 
-def _get_latest_values(existing_content, scope=None):
+def get_latest_values(existing_content, scope=None):
     """Get the latest "registries", "py_versions" and "repository" values
 
     Args:
@@ -40,10 +40,10 @@ def _get_latest_values(existing_content, scope=None):
                 )
             )
 
-    latest_version = list(existing_content[scope]["versions"].keys())[-1]
-    registries = existing_content[scope]["versions"][latest_version]["registries"]
-    py_versions = existing_content[scope]["versions"][latest_version]["py_versions"]
-    repository = existing_content[scope]["versions"][latest_version]["repository"]
+    latest_version = list(existing_content["versions"].keys())[-1]
+    registries = existing_content["versions"][latest_version]["registries"]
+    py_versions = existing_content["versions"][latest_version]["py_versions"]
+    repository = existing_content["versions"][latest_version]["repository"]
 
     return registries, py_versions, repository
 
@@ -173,7 +173,7 @@ def add_version(
      """
     py_versions = py_versions.split(",")
     processors = processors.split(",")
-    latest_registries, latest_py_versions, latest_repository = _get_latest_values(
+    latest_registries, latest_py_versions, latest_repository = get_latest_values(
         existing_content, scope
     )
     if not py_versions:

tests/unit/sagemaker/cli/test_framework_version_upgrade_tool.py renamed to tests/unit/sagemaker/cli/test_framework_upgrade_tool.py

@@ -218,3 +218,22 @@ def test_algo_add_region(algo_region_content, algo_expected_region_content):
     account = "987654321098"
     framework_upgrade.add_region(algo_region_content, region, account)
     assert algo_region_content == algo_expected_region_content
+
+
+def test_dlc_get_latest_content(dlc_content):
+    latest_version = "1.0.0"
+    scope = "eia"
+    registries, py_versions, repository = framework_upgrade.get_latest_values(
+        dlc_content, scope=scope
+    )
+    assert registries == dlc_content[scope]["versions"][latest_version]["registries"]
+    assert py_versions == dlc_content[scope]["versions"][latest_version]["py_versions"]
+    assert repository == dlc_content[scope]["versions"][latest_version]["repository"]
+
+
+def test_algo_get_latest_content(algo_content):
+    latest_version = "0.10"
+    registries, py_versions, repository = framework_upgrade.get_latest_values(algo_content)
+    assert registries == algo_content["versions"][latest_version]["registries"]
+    assert py_versions == algo_content["versions"][latest_version]["py_versions"]
+    assert repository == algo_content["versions"][latest_version]["repository"]