[MIGRATE] Fix edge case bugs with az migrate local replication init & new commands (#9675)

* Create extension

* Update src/migrate/azext_migrate/__init__.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Fix import issues

* Update src/migrate/setup.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Small

* Small lint

* Small

* disable lint for this check

* Add json

* Fix licesnse issue

* fix small

* Small

* Get rid of unused variables

* Add service name and code owner

* New version

* Style

* Small

* Update

* Follow standard

* Add suggestions

* Small

* Not preview

* Add flag to become experimental

* Update history

* Fix

* small

* Create get job and remove replication commands

* Revert "Create get job and remove replication commands"

This reverts commit 0f7acb7.

* Update version

* Sync with other branch

* Fix error in init command

* Remove duplicate files

* Fix local issues

* Select correct fabric

* Process project name

---------

Co-authored-by: Saif Al-Din Ali <saifaldinali@microsoft.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

src/migrate/HISTORY.rst

@@ -2,6 +2,9 @@
 
 Release History
 ===============
+3.0.0b4
++++++++++++++++
+* Fix edge case bugs in az migrate local replication init & new commands.
 
 3.0.0b3
 +++++++++++++++

src/migrate/azext_migrate/custom.py

@@ -172,7 +172,7 @@ def new_local_server_replication(cmd,
         subscription_id = get_subscription_id(cmd.cli_ctx)
         print(f"Selected Subscription Id: '{subscription_id}'")
 
-    rg_uri, machine_id, subscription_id = validate_server_parameters(
+    rg_uri, machine_id, subscription_id, project_name = validate_server_parameters(
         cmd,
         machine_id,
         machine_index,
@@ -239,7 +239,7 @@ def new_local_server_replication(cmd,
                 "Please verify your appliance setup and provided "
                 "-machine_id.")
 
-        amh_solution, migrate_project, machine_props = process_amh_solution(
+        amh_solution, migrate_project, machine_props, project_name = process_amh_solution(
             cmd,
             machine,
             site_object,

src/migrate/azext_migrate/helpers/replication/init/_setup_extension.py

@@ -260,7 +260,8 @@ def build_extension_body(instance_type, source_fabric_id,
 
 
 def _wait_for_extension_creation(cmd, extension_uri):
-    """Wait for extension creation to complete."""
+    """Wait for extension creation to complete. Returns final state."""
+    ext_state = None
     for i in range(20):
         time.sleep(30)
         try:
@@ -277,6 +278,7 @@ def _wait_for_extension_creation(cmd, extension_uri):
                     break
         except CLIError:
             print(f"Waiting for extension... ({i + 1}/20)")
+    return ext_state
 
 
 def _handle_extension_creation_error(cmd, extension_uri, create_error):
@@ -315,7 +317,20 @@ def create_replication_extension(cmd, extension_uri, extension_body):
             print("Extension creation initiated successfully")
             # Wait for the extension to be created
             print("Waiting for extension creation to complete...")
-            _wait_for_extension_creation(cmd, extension_uri)
+            ext_state = _wait_for_extension_creation(cmd, extension_uri)
+            if ext_state == ProvisioningState.Failed.value:
+                raise CLIError(
+                    "Replication extension creation failed. "
+                    "Check the extension resource in the Azure portal "
+                    "for detailed error information.")
+            if ext_state == ProvisioningState.Canceled.value:
+                raise CLIError(
+                    "Replication extension creation was canceled.")
+            if ext_state is None:
+                raise CLIError(
+                    "Replication extension creation timed out after "
+                    "10 minutes. Check the extension status in the "
+                    "Azure portal.")
     except CLIError as create_error:
         _handle_extension_creation_error(cmd, extension_uri, create_error)
 

src/migrate/azext_migrate/helpers/replication/init/_setup_policy.py

@@ -101,13 +101,22 @@ def find_fabric(all_fabrics, appliance_name, fabric_instance_type,
             })
 
         if is_succeeded and is_correct_instance and name_matches:
-            # If solution doesn't match, log warning but still consider it
-            if not is_correct_solution:
-                logger.warning(
-                    "Fabric '%s' matches name and type but has "
-                    "different solution ID", fabric_name)
-            fabric = candidate
-            break
+            if is_correct_solution:
+                # Perfect match - use it immediately
+                fabric = candidate
+                break
+            # Name/type match but wrong solution ID - keep as fallback
+            if not fabric:
+                fabric = candidate
+
+    if fabric:
+        fabric_props = fabric.get('properties', {}).get('customProperties', {})
+        fabric_sol_id = fabric_props.get('migrationSolutionId', '').rstrip('/')
+        expected_sol_id = amh_solution.get('id', '').rstrip('/')
+        if fabric_sol_id.lower() != expected_sol_id.lower():
+            logger.warning(
+                "Fabric '%s' matches name and type but has "
+                "different solution ID", fabric.get('name'))
 
     if not fabric:
         appliance_type_label = "source" if is_source else "target"
@@ -170,6 +179,7 @@ def find_fabric(all_fabrics, appliance_name, fabric_instance_type,
 def get_fabric_agent(cmd, replication_fabrics_uri, fabric, appliance_name,
                      fabric_instance_type):
     """Get and validate fabric agent (DRA) for the given fabric."""
+    logger = get_logger(__name__)
     fabric_name = fabric.get('name')
     dras_uri = (
         f"{replication_fabrics_uri}/{fabric_name}"
@@ -180,18 +190,59 @@ def get_fabric_agent(cmd, replication_fabrics_uri, fabric, appliance_name,
     dras = dras_response.json().get('value', [])
 
     dra = None
+    found_but_not_responsive = None
     for candidate in dras:
         props = candidate.get('properties', {})
         custom_props = props.get('customProperties', {})
-        if (props.get('machineName') == appliance_name and
-                custom_props.get('instanceType') == fabric_instance_type and
-                bool(props.get('isResponsive'))):
-            dra = candidate
-            break
+        machine_name = props.get('machineName', '')
+        if (machine_name.lower() == appliance_name.lower() and
+                custom_props.get('instanceType') == fabric_instance_type):
+            if bool(props.get('isResponsive')):
+                dra = candidate
+                break
+            found_but_not_responsive = candidate
+
+    # Accept a non-responsive DRA if it's the only match and is provisioned
+    if not dra and found_but_not_responsive:
+        nr_props = found_but_not_responsive.get('properties', {})
+        last_heartbeat = nr_props.get('lastHeartbeat', 'unknown')
+        if (nr_props.get('provisioningState') ==
+                ProvisioningState.Succeeded.value):
+            logger.warning(
+                "The appliance '%s' DRA is not responsive "
+                "(last heartbeat: %s). Proceeding since provisioning "
+                "state is 'Succeeded'.",
+                appliance_name, last_heartbeat)
+            dra = found_but_not_responsive
+        else:
+            raise CLIError(
+                f"The appliance '{appliance_name}' is in a "
+                f"disconnected state (last heartbeat: {last_heartbeat}, "
+                f"provisioningState: "
+                f"{nr_props.get('provisioningState')})."
+            )
 
     if not dra:
+        # Log available DRAs for diagnostics
+        if dras:
+            logger.warning(
+                "No matching fabric agent found for appliance '%s' "
+                "(expected instanceType '%s'). Available agents:",
+                appliance_name, fabric_instance_type)
+            for candidate in dras:
+                props = candidate.get('properties', {})
+                custom_props = props.get('customProperties', {})
+                logger.warning(
+                    "  - machineName: '%s', instanceType: '%s', "
+                    "isResponsive: %s",
+                    props.get('machineName'),
+                    custom_props.get('instanceType'),
+                    props.get('isResponsive'))
+
         raise CLIError(
-            f"The appliance '{appliance_name}' is in a disconnected state."
+            f"No fabric agent found for appliance '{appliance_name}' "
+            f"on fabric '{fabric_name}'. Verify that the appliance is "
+            f"properly registered and connected."
         )
 
     return dra

src/migrate/azext_migrate/helpers/replication/new/_process_inputs.py

@@ -228,7 +228,7 @@ def process_amh_solution(cmd,
             f"'{resource_group_name}' and project '{project_name}'. "
             "Please verify your appliance setup."
         )
-    return amh_solution, migrate_project, machine_props
+    return amh_solution, migrate_project, machine_props, project_name
 
 
 def process_replication_vault(cmd,
@@ -482,15 +482,23 @@ def _process_source_fabrics(all_fabrics,
             })
 
         if is_succeeded and is_correct_instance and name_matches:
-            # If solution doesn't match, log warning but still consider it
-            if not is_correct_solution:
-                logger.warning(
-                    "Fabric '%s' matches name and type but has different "
-                    "solution ID",
-                    fabric_name
-                )
-            source_fabric = fabric
-            break
+            if is_correct_solution:
+                source_fabric = fabric
+                break
+            if not source_fabric:
+                source_fabric = fabric
+
+    if source_fabric:
+        sf_props = source_fabric.get('properties', {}).get(
+            'customProperties', {})
+        sf_sol_id = sf_props.get('migrationSolutionId', '').rstrip('/')
+        exp_sol_id = amh_solution.get('id', '').rstrip('/')
+        if sf_sol_id.lower() != exp_sol_id.lower():
+            logger.warning(
+                "Fabric '%s' matches name and type but has different "
+                "solution ID",
+                source_fabric.get('name'))
+
     return source_fabric, source_fabric_candidates
 
 
@@ -679,12 +687,22 @@ def _process_target_fabrics(all_fabrics,
             })
 
         if is_succeeded and is_correct_instance and name_matches:
-            if not is_correct_solution:
-                logger.warning(
-                    "Fabric '%s' matches name and type but has different "
-                    "solution ID", fabric_name)
-            target_fabric = fabric
-            break
+            if is_correct_solution:
+                target_fabric = fabric
+                break
+            if not target_fabric:
+                target_fabric = fabric
+
+    if target_fabric:
+        tf_props = target_fabric.get('properties', {}).get(
+            'customProperties', {})
+        tf_sol_id = tf_props.get('migrationSolutionId', '').rstrip('/')
+        exp_sol_id = amh_solution.get('id', '').rstrip('/')
+        if tf_sol_id.lower() != exp_sol_id.lower():
+            logger.warning(
+                "Fabric '%s' matches name and type but has different "
+                "solution ID", target_fabric.get('name'))
+
     return target_fabric, target_fabric_candidates, \
         target_fabric_instance_type
 
@@ -731,28 +749,48 @@ def process_target_fabric(cmd,
                           amh_solution):
     # Get source fabric agent (DRA)
     source_fabric_name = source_fabric.get('name')
-    dras_uri = (
+    source_dras = send_get_request(
+        cmd,
         f"{rg_uri}/providers/Microsoft.DataReplication"
         f"/replicationFabrics/{source_fabric_name}/fabricAgents"
         f"?api-version={APIVersion.Microsoft_DataReplication.value}"
-    )
-    source_dras_response = send_get_request(cmd, dras_uri)
-    source_dras = source_dras_response.json().get('value', [])
+    ).json().get('value', [])
 
     source_dra = None
+    source_found_not_responsive = None
     for dra in source_dras:
         props = dra.get('properties', {})
         custom_props = props.get('customProperties', {})
-        if (props.get('machineName') == source_appliance_name and
-                custom_props.get('instanceType') == fabric_instance_type and
-                bool(props.get('isResponsive'))):
-            source_dra = dra
-            break
+        machine_name = props.get('machineName', '')
+        if (machine_name.lower() == source_appliance_name.lower() and
+                custom_props.get('instanceType') == fabric_instance_type):
+            if bool(props.get('isResponsive')):
+                source_dra = dra
+                break
+            source_found_not_responsive = dra
+
+    if not source_dra and source_found_not_responsive:
+        nr_props = source_found_not_responsive.get('properties', {})
+        last_hb = nr_props.get('lastHeartbeat', 'unknown')
+        if (nr_props.get('provisioningState') ==
+                ProvisioningState.Succeeded.value):
+            logger.warning(
+                "The source appliance '%s' DRA is not responsive "
+                "(last heartbeat: %s). Proceeding since provisioning "
+                "state is 'Succeeded'.",
+                source_appliance_name, last_hb)
+            source_dra = source_found_not_responsive
+        else:
+            raise CLIError(
+                f"The source appliance '{source_appliance_name}' is in a "
+                f"disconnected state (last heartbeat: {last_hb}).")
 
     if not source_dra:
         raise CLIError(
-            f"The source appliance '{source_appliance_name}' is in a "
-            f"disconnected state.")
+            f"No fabric agent found for source appliance "
+            f"'{source_appliance_name}' on fabric "
+            f"'{source_fabric_name}'. Verify that the appliance is "
+            f"properly registered and connected.")
 
     target_fabric, target_fabric_candidates, \
         target_fabric_instance_type = _process_target_fabrics(
@@ -769,28 +807,48 @@ def process_target_fabric(cmd,
 
     # Get target fabric agent (DRA)
     target_fabric_name = target_fabric.get('name')
-    target_dras_uri = (
+    target_dras = send_get_request(
+        cmd,
         f"{rg_uri}/providers/Microsoft.DataReplication"
         f"/replicationFabrics/{target_fabric_name}/fabricAgents"
         f"?api-version={APIVersion.Microsoft_DataReplication.value}"
-    )
-    target_dras_response = send_get_request(cmd, target_dras_uri)
-    target_dras = target_dras_response.json().get('value', [])
+    ).json().get('value', [])
 
     target_dra = None
+    target_found_not_responsive = None
     for dra in target_dras:
         props = dra.get('properties', {})
         custom_props = props.get('customProperties', {})
-        if (props.get('machineName') == target_appliance_name and
+        machine_name = props.get('machineName', '')
+        if (machine_name.lower() == target_appliance_name.lower() and
                 custom_props.get('instanceType') ==
-                target_fabric_instance_type and
-                bool(props.get('isResponsive'))):
-            target_dra = dra
-            break
+                target_fabric_instance_type):
+            if bool(props.get('isResponsive')):
+                target_dra = dra
+                break
+            target_found_not_responsive = dra
+
+    if not target_dra and target_found_not_responsive:
+        nr_props = target_found_not_responsive.get('properties', {})
+        last_hb = nr_props.get('lastHeartbeat', 'unknown')
+        if (nr_props.get('provisioningState') ==
+                ProvisioningState.Succeeded.value):
+            logger.warning(
+                "The target appliance '%s' DRA is not responsive "
+                "(last heartbeat: %s). Proceeding since provisioning "
+                "state is 'Succeeded'.",
+                target_appliance_name, last_hb)
+            target_dra = target_found_not_responsive
+        else:
+            raise CLIError(
+                f"The target appliance '{target_appliance_name}' is in a "
+                f"disconnected state (last heartbeat: {last_hb}).")
 
     if not target_dra:
         raise CLIError(
-            f"The target appliance '{target_appliance_name}' is in a "
-            f"disconnected state.")
+            f"No fabric agent found for target appliance "
+            f"'{target_appliance_name}' on fabric "
+            f"'{target_fabric_name}'. Verify that the appliance is "
+            f"properly registered and connected.")
 
     return target_fabric, source_dra, target_dra

src/migrate/azext_migrate/helpers/replication/new/_validate.py

@@ -206,6 +206,11 @@ def validate_server_parameters(  # pylint: disable=too-many-locals,too-many-bran
         # machine_id was provided directly
         # Check if it's in Microsoft.Migrate format and needs to be resolved
         if "/Microsoft.Migrate/MigrateProjects/" in machine_id or "/Microsoft.Migrate/migrateprojects/" in machine_id:
+            # Extract project_name from the Microsoft.Migrate machine ID
+            migrate_id_parts = machine_id.split("/")
+            if len(migrate_id_parts) >= 9 and not project_name:
+                project_name = migrate_id_parts[8]
+
             # This is a Migrate Project machine ID, need to resolve to OffAzure machine ID
             migrate_machine = get_resource_by_id(
                 cmd, machine_id, APIVersion.Microsoft_Migrate.value)
@@ -257,7 +262,7 @@ def validate_server_parameters(  # pylint: disable=too-many-locals,too-many-bran
             f"/subscriptions/{subscription_id}/"
             f"resourceGroups/{resource_group_name}")
 
-    return rg_uri, machine_id, subscription_id
+    return rg_uri, machine_id, subscription_id, project_name
 
 
 def validate_required_parameters(machine_id,