diff --git a/src/connectedk8s/azext_connectedk8s/_precheckutils.py b/src/connectedk8s/azext_connectedk8s/_precheckutils.py index b10c9a22241..231f2b4c659 100644 --- a/src/connectedk8s/azext_connectedk8s/_precheckutils.py +++ b/src/connectedk8s/azext_connectedk8s/_precheckutils.py @@ -212,7 +212,7 @@ def executing_cluster_diagnostic_checks_job( ) return None - mcr_url = azext_utils.get_mcr_path(cmd) + mcr_url = azext_utils.get_mcr_path(cmd.cli_ctx.cloud.endpoints.active_directory) chart_path = azext_utils.get_chart_path( f"{mcr_url}/{consts.Cluster_Diagnostic_Checks_Job_Registry_Path}", diff --git a/src/connectedk8s/azext_connectedk8s/_utils.py b/src/connectedk8s/azext_connectedk8s/_utils.py index 6717ed907c1..e3268806f7b 100644 --- a/src/connectedk8s/azext_connectedk8s/_utils.py +++ b/src/connectedk8s/azext_connectedk8s/_utils.py @@ -58,11 +58,15 @@ # pylint: disable=bare-except -def get_mcr_path(cmd: CLICommand) -> str: - active_directory_array = cmd.cli_ctx.cloud.endpoints.active_directory.split(".") +def get_mcr_path(active_directory_endpoint: str) -> str: + active_directory_array = active_directory_endpoint.split(".") - # default for public, mc, ff clouds - mcr_postfix = active_directory_array[2] + # For US Government and China clouds, use public mcr + if active_directory_endpoint.endswith((".us", ".cn")): + return "mcr.microsoft.com" + + # Default MCR postfix + mcr_postfix = "com" # special cases for USSec, exclude part of suffix if len(active_directory_array) == 4 and active_directory_array[2] == "microsoft": mcr_postfix = active_directory_array[3] diff --git a/src/connectedk8s/azext_connectedk8s/clientproxyhelper/_binaryutils.py b/src/connectedk8s/azext_connectedk8s/clientproxyhelper/_binaryutils.py index c655b4269de..5097fb952cc 100644 --- a/src/connectedk8s/azext_connectedk8s/clientproxyhelper/_binaryutils.py +++ b/src/connectedk8s/azext_connectedk8s/clientproxyhelper/_binaryutils.py @@ -76,7 +76,7 @@ def _download_proxy_from_MCR( operating_system: str, architecture: str, ) -> None: - mcr_url = utils.get_mcr_path(cmd) + mcr_url = utils.get_mcr_path(cmd.cli_ctx.cloud.endpoints.active_directory) mar_target = f"{mcr_url}/{consts.CLIENT_PROXY_MCR_TARGET}/{operating_system.lower()}/{architecture}/arc-proxy" logger.debug( diff --git a/src/connectedk8s/azext_connectedk8s/custom.py b/src/connectedk8s/azext_connectedk8s/custom.py index c5ed8ed8ec9..d90c0c760c0 100644 --- a/src/connectedk8s/azext_connectedk8s/custom.py +++ b/src/connectedk8s/azext_connectedk8s/custom.py @@ -1318,7 +1318,7 @@ def install_helm_client(cmd: CLICommand) -> str: "Downloading helm client for first time. This can take few minutes..." ) - mcr_url = utils.get_mcr_path(cmd) + mcr_url = utils.get_mcr_path(cmd.cli_ctx.cloud.endpoints.active_directory) client = oras.client.OrasClient(hostname=mcr_url) retry_count = 3 diff --git a/src/connectedk8s/azext_connectedk8s/tests/unittests/test_utils_.py b/src/connectedk8s/azext_connectedk8s/tests/unittests/test_utils_.py index 768290b902a..32d1da1e3b4 100644 --- a/src/connectedk8s/azext_connectedk8s/tests/unittests/test_utils_.py +++ b/src/connectedk8s/azext_connectedk8s/tests/unittests/test_utils_.py @@ -9,6 +9,7 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))) from azext_connectedk8s._utils import ( + get_mcr_path, process_helm_error_detail, redact_sensitive_fields_from_string, remove_rsa_private_key, @@ -76,5 +77,27 @@ def test_redact_sensitive_fields_from_string(): ) +def test_get_mcr_path(): + input_active_directory = "login.microsoftonline.com" + expected_output = "mcr.microsoft.com" + assert get_mcr_path(input_active_directory) == expected_output + + input_active_directory = "login.microsoftonline.us" + expected_output = "mcr.microsoft.com" + assert get_mcr_path(input_active_directory) == expected_output + + input_active_directory = "login.chinacloudapi.cn" + expected_output = "mcr.microsoft.com" + assert get_mcr_path(input_active_directory) == expected_output + + input_active_directory = "https://login.microsoftonline.microsoft.foo" + expected_output = "mcr.microsoft.foo" + assert get_mcr_path(input_active_directory) == expected_output + + input_active_directory = "https://login.microsoftonline.some.cloud.bar" + expected_output = "mcr.microsoft.some.cloud.bar" + assert get_mcr_path(input_active_directory) == expected_output + + if __name__ == "__main__": pytest.main() diff --git a/src/k8s-extension/azext_k8s_extension/custom.py b/src/k8s-extension/azext_k8s_extension/custom.py index f7626d06e58..6122ea61224 100644 --- a/src/k8s-extension/azext_k8s_extension/custom.py +++ b/src/k8s-extension/azext_k8s_extension/custom.py @@ -763,7 +763,7 @@ def install_helm_client(cmd: CLICommand) -> str: "Downloading helm client for first time. This can take few minutes..." ) - mcr_url = utils.get_mcr_path(cmd) + mcr_url = utils.get_mcr_path(cmd.cli_ctx.cloud.endpoints.active_directory) client = oras.client.OrasClient(hostname=mcr_url) retry_count = 3 diff --git a/src/k8s-extension/azext_k8s_extension/utils.py b/src/k8s-extension/azext_k8s_extension/utils.py index 767550ccef8..345d92f9016 100644 --- a/src/k8s-extension/azext_k8s_extension/utils.py +++ b/src/k8s-extension/azext_k8s_extension/utils.py @@ -361,11 +361,15 @@ def create_folder_diagnosticlogs(folder_name: str, base_folder_name: str) -> tup ) return "", False -def get_mcr_path(cmd: CLICommand) -> str: - active_directory_array = cmd.cli_ctx.cloud.endpoints.active_directory.split(".") +def get_mcr_path(active_directory_endpoint: str) -> str: + active_directory_array = active_directory_endpoint.split(".") - # default for public, mc, ff clouds - mcr_postfix = active_directory_array[2] + # For US Government and China clouds, use public mcr + if active_directory_endpoint.endswith((".us", ".cn")): + return "mcr.microsoft.com" + + # Default MCR postfix + mcr_postfix = "com" # special cases for USSec, exclude part of suffix if len(active_directory_array) == 4 and active_directory_array[2] == "microsoft": mcr_postfix = active_directory_array[3] diff --git a/testing/test/configurations/EnableDisableFeatures.Tests.ps1 b/testing/test/configurations/EnableDisableFeatures.Tests.ps1 index ce5ff9bd367..8c6f1f5d1f3 100644 --- a/testing/test/configurations/EnableDisableFeatures.Tests.ps1 +++ b/testing/test/configurations/EnableDisableFeatures.Tests.ps1 @@ -106,6 +106,6 @@ Describe 'ConnectedK8s Enable Disable Features Scenario' { # Configuration should be removed from the resource model - expect ResourceNotFound error $output = Invoke-AzCommand "az connectedk8s show -n $($ENVCONFIG.arcClusterName) -g $($ENVCONFIG.resourceGroup)" 2>&1 - $output | Should -Match "ResourceNotFound" + $output | Should -Match "(ResourceNotFound|could not be found|not found)" } } \ No newline at end of file