Bug Fix for FFX mcr url

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}",

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]

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(

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

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()

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

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]

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)"
     }
 }