Add unit test for _get_codeartifact_index

Author: akuma12

src/sagemaker/processing.py

@@ -22,7 +22,7 @@
 import pathlib
 import logging
 from textwrap import dedent
-from typing import Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 from copy import copy
 import re
 
@@ -1846,14 +1846,18 @@ def _pack_and_upload_code(
 
         return s3_runproc_sh, inputs, job_name
 
-    def _get_codeartifact_index(self, codeartifact_repo_arn: str):
+    def _get_codeartifact_index(self, codeartifact_repo_arn: str, codeartifact_client: Any = None):
         """
         Build the authenticated codeartifact index url based on the arn provided
         via codeartifact_repo_arn property following the form
         # `arn:${Partition}:codeartifact:${Region}:${Account}:repository/${Domain}/${Repository}`
         https://docs.aws.amazon.com/codeartifact/latest/ug/python-configure-pip.html
         https://docs.aws.amazon.com/service-authorization/latest/reference/list_awscodeartifact.html#awscodeartifact-resources-for-iam-policies
-        :return: authenticated codeartifact index url
+        Args:
+            codeartifact_repo_arn: arn of the codeartifact repository
+            codeartifact_client: boto3 client for codeartifact (used for testing)
+        Returns:
+            authenticated codeartifact index url
         """
 
         arn_regex = (
@@ -1862,7 +1866,7 @@ def _get_codeartifact_index(self, codeartifact_repo_arn: str):
         )
         m = re.match(arn_regex, codeartifact_repo_arn)
         if not m:
-            raise Exception("invalid CodeArtifact repository arn {}".format(codeartifact_repo_arn))
+            raise ValueError("invalid CodeArtifact repository arn {}".format(codeartifact_repo_arn))
         domain = m.group("domain")
         owner = m.group("account")
         repository = m.group("repository")
@@ -1877,10 +1881,12 @@ def _get_codeartifact_index(self, codeartifact_repo_arn: str):
             region,
         )
         try:
-            client = self.sagemaker_session.boto_session.client("codeartifact", region_name=region)
-            auth_token_response = client.get_authorization_token(domain=domain, domainOwner=owner)
+            if not codeartifact_client:
+                codeartifact_client = self.sagemaker_session.boto_session.client("codeartifact", region_name=region)
+            
+            auth_token_response = codeartifact_client.get_authorization_token(domain=domain, domainOwner=owner)
             token = auth_token_response["authorizationToken"]
-            endpoint_response = client.get_repository_endpoint(
+            endpoint_response = codeartifact_client.get_repository_endpoint(
                 domain=domain, domainOwner=owner, repository=repository, format="pypi"
             )
             unauthenticated_index = endpoint_response["repositoryEndpoint"]
@@ -1893,9 +1899,9 @@ def _get_codeartifact_index(self, codeartifact_repo_arn: str):
                     unauthenticated_index,
                 ),
             )
-        except Exception:
-            logger.error("failed to configure pip to use codeartifact")
-            raise Exception("failed to configure pip to use codeartifact")
+        except Exception as e:
+            logger.error("failed to configure pip to use codeartifact: %s", e, exc_info=True)
+            raise RuntimeError("failed to configure pip to use codeartifact")
 
     def _generate_framework_script(
         self, user_script: str, codeartifact_repo_arn: str = None

tests/unit/test_processing.py

@@ -13,7 +13,10 @@
 from __future__ import absolute_import
 
 import copy
+import datetime
 
+import boto3
+from botocore.stub import Stubber
 import pytest
 from mock import Mock, patch, MagicMock
 from packaging import version
@@ -1099,6 +1102,121 @@ def test_pyspark_processor_configuration_path_pipeline_config(
     )
 
 
+@patch("sagemaker.workflow.utilities._pipeline_config", MOCKED_PIPELINE_CONFIG)
+def test_get_codeartifact_index(pipeline_session):
+    codeartifact_repo_arn = "arn:aws:codeartifact:us-west-2:012345678901:repository/test-domain/test-repository"
+    codeartifact_url = "test-domain-012345678901.d.codeartifact.us-west-2.amazonaws.com/pypi/test-repository/simple/"
+
+    client = boto3.client('codeartifact', region_name=REGION)
+    stubber = Stubber(client)
+        
+    get_auth_token_response = {
+        "authorizationToken": "mocked_token",
+        "expiration": datetime.datetime(2045, 1, 1, 0, 0, 0)
+    }
+    auth_token_expected_params = {"domain": "test-domain", "domainOwner": "012345678901"}
+    stubber.add_response("get_authorization_token", get_auth_token_response, auth_token_expected_params)
+
+    get_repo_endpoint_response = {"repositoryEndpoint": f"https://{codeartifact_url}"}
+    repo_endpoint_expected_params = {
+        "domain": "test-domain",
+        "domainOwner": "012345678901",
+        "repository": "test-repository",
+        "format": "pypi"
+    }
+    stubber.add_response("get_repository_endpoint", get_repo_endpoint_response, repo_endpoint_expected_params)
+
+    processor = PyTorchProcessor(
+        role=ROLE,
+        instance_type="ml.m4.xlarge",
+        framework_version="2.0.1",
+        py_version="py310",
+        instance_count=1,
+        sagemaker_session=pipeline_session,
+    )
+
+    with stubber:
+        codeartifact_index = processor._get_codeartifact_index(codeartifact_repo_arn=codeartifact_repo_arn, codeartifact_client=client)
+    
+    assert codeartifact_index == f"https://aws:mocked_token@{codeartifact_url}"
+
+
+@patch("sagemaker.workflow.utilities._pipeline_config", MOCKED_PIPELINE_CONFIG)
+def test_get_codeartifact_index_bad_repo_arn(pipeline_session):
+    codeartifact_repo_arn = "arn:aws:codeartifact:us-west-2:012345678901:repository/test-domain"
+    codeartifact_url = "test-domain-012345678901.d.codeartifact.us-west-2.amazonaws.com/pypi/test-repository/simple/"
+
+    client = boto3.client('codeartifact', region_name=REGION)
+    stubber = Stubber(client)
+        
+    get_auth_token_response = {
+        "authorizationToken": "mocked_token",
+        "expiration": datetime.datetime(2045, 1, 1, 0, 0, 0)
+    }
+    auth_token_expected_params = {"domain": "test-domain", "domainOwner": "012345678901"}
+    stubber.add_response("get_authorization_token", get_auth_token_response, auth_token_expected_params)
+
+    get_repo_endpoint_response = {"repositoryEndpoint": f"https://{codeartifact_url}"}
+    repo_endpoint_expected_params = {
+        "domain": "test-domain",
+        "domainOwner": "012345678901",
+        "repository": "test-repository",
+        "format": "pypi"
+    }
+    stubber.add_response("get_repository_endpoint", get_repo_endpoint_response, repo_endpoint_expected_params)
+
+    processor = PyTorchProcessor(
+        role=ROLE,
+        instance_type="ml.m4.xlarge",
+        framework_version="2.0.1",
+        py_version="py310",
+        instance_count=1,
+        sagemaker_session=pipeline_session,
+    )
+
+    with stubber:
+        with pytest.raises(ValueError):
+            processor._get_codeartifact_index(codeartifact_repo_arn=codeartifact_repo_arn, codeartifact_client=client)
+
+
+@patch("sagemaker.workflow.utilities._pipeline_config", MOCKED_PIPELINE_CONFIG)
+def test_get_codeartifact_index_client_error(pipeline_session):
+    codeartifact_repo_arn = "arn:aws:codeartifact:us-west-2:012345678901:repository/test-domain/test-repository"
+    codeartifact_url = "test-domain-012345678901.d.codeartifact.us-west-2.amazonaws.com/pypi/test-repository/simple/"
+
+    client = boto3.client('codeartifact', region_name=REGION)
+    stubber = Stubber(client)
+        
+    get_auth_token_response = {
+        "authorizationToken": "mocked_token",
+        "expiration": datetime.datetime(2045, 1, 1, 0, 0, 0)
+    }
+    auth_token_expected_params = {"domain": "test-domain", "domainOwner": "012345678901"}
+    stubber.add_client_error("get_authorization_token", service_error_code="404", expected_params=auth_token_expected_params)
+
+    get_repo_endpoint_response = {"repositoryEndpoint": f"https://{codeartifact_url}"}
+    repo_endpoint_expected_params = {
+        "domain": "test-domain",
+        "domainOwner": "012345678901",
+        "repository": "test-repository",
+        "format": "pypi"
+    }
+    stubber.add_response("get_repository_endpoint", get_repo_endpoint_response, repo_endpoint_expected_params)
+
+    processor = PyTorchProcessor(
+        role=ROLE,
+        instance_type="ml.m4.xlarge",
+        framework_version="2.0.1",
+        py_version="py310",
+        instance_count=1,
+        sagemaker_session=pipeline_session,
+    )
+
+    with stubber:
+        with pytest.raises(RuntimeError):
+            processor._get_codeartifact_index(codeartifact_repo_arn=codeartifact_repo_arn, codeartifact_client=client)
+
+
 def _get_script_processor(sagemaker_session):
     return ScriptProcessor(
         role=ROLE,