From 532951a0e1c423c84b5cd33d356e6c368e1239b5 Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Tue, 29 Jul 2025 22:12:13 +0000
Subject: [PATCH 01/16] initial changes for stack-hci vmconnect command

---
 src/stack-hci/azext_stack_hci/commands.py |   5 +
 src/stack-hci/azext_stack_hci/custom.py   | 134 ++++++++++++++++++++++
 2 files changed, 139 insertions(+)

diff --git a/src/stack-hci/azext_stack_hci/commands.py b/src/stack-hci/azext_stack_hci/commands.py
index a5f9cf4ff53..fcf39640117 100644
--- a/src/stack-hci/azext_stack_hci/commands.py
+++ b/src/stack-hci/azext_stack_hci/commands.py
@@ -19,3 +19,8 @@ def load_command_table(self, _):  # pylint: disable=unused-argument
         self.command_table['stack-hci cluster create'] = ClusterCreate(loader=self)
         self.command_table['stack-hci cluster identity assign'] = IdentityAssign(loader=self)
         self.command_table['stack-hci cluster identity remove'] = IdentityRemove(loader=self)
+
+    with self.command_group('stack-hci vmconnect'):
+        from azext_stack_hci.custom import VmConnectEnable, VmConnectDisable
+        self.command_table['stack-hci vmconnect enable'] = VmConnectEnable(loader=self)
+        self.command_table['stack-hci vmconnect disable'] = VmConnectDisable(loader=self)
diff --git a/src/stack-hci/azext_stack_hci/custom.py b/src/stack-hci/azext_stack_hci/custom.py
index c59c1778af3..c62b891bdbf 100644
--- a/src/stack-hci/azext_stack_hci/custom.py
+++ b/src/stack-hci/azext_stack_hci/custom.py
@@ -75,3 +75,137 @@ def pre_instance_update(self, instance):
         args = self.ctx.args
         if args.system_assigned:
             args.type = 'None'
+
+
+class VmConnectEnable:
+    """Enable VM Connect functionality for Stack HCI cluster."""
+    
+    def __init__(self, loader=None):
+        self.loader = loader
+    
+    @classmethod
+    def _build_arguments_schema(cls, *args, **kwargs):
+        from azure.cli.core.aaz import AAZResourceGroupNameArg, AAZStrArg, AAZDictArg
+        args_schema = {}
+        args_schema.cluster_name = AAZStrArg(
+            options=["--cluster-name", "-n"],
+            help="The name of the cluster.",
+            required=True
+        )
+        args_schema.resource_group = AAZResourceGroupNameArg(
+            options=["--resource-group", "-g"],
+            help="Name of resource group.",
+            required=True
+        )
+        args_schema.vm_name = AAZStrArg(
+            options=["--vm-name"],
+            help="The name of the virtual machine.",
+            required=True
+        )
+        return args_schema
+    
+    def __call__(self, cmd, **kwargs):
+        from azure.cli.core.commands.client_factory import get_subscription_id
+        from azure.cli.core.util import send_raw_request
+        import json
+        
+        cluster_name = kwargs.get('cluster_name')
+        resource_group = kwargs.get('resource_group')
+        vm_name = kwargs.get('vm_name')
+        
+        subscription_id = get_subscription_id(cmd.cli_ctx)
+        
+        # Construct the REST API path
+        path = f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers/Microsoft.AzureStackHCI/clusters/{cluster_name}/jobs/VmConnectProvision"
+        
+        # API version
+        api_version = "2023-12-01-preview"
+        url = f"https://management.azure.com{path}?api-version={api_version}"
+        
+        # Default payload with VM name
+        payload = {
+            "properties": {
+                "jobType": "VmConnectProvision",
+                "deploymentMode": "Deploy",
+                "vmConnectProvisionJobDetails": [
+                    {
+                        "vmName": vm_name
+                    }
+                ]
+            }
+        }
+        
+        # Make the REST API call
+        try:
+            response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
+            return response.json() if response.content else {"message": f"VM Connect provision job initiated successfully for VM: {vm_name}"}
+        except Exception as e:
+            from azure.cli.core.util import CLIError
+            raise CLIError(f"Failed to enable VM Connect for VM '{vm_name}': {str(e)}")
+
+
+class VmConnectDisable:
+    """Disable VM Connect functionality for Stack HCI cluster."""
+    
+    def __init__(self, loader=None):
+        self.loader = loader
+    
+    @classmethod
+    def _build_arguments_schema(cls, *args, **kwargs):
+        from azure.cli.core.aaz import AAZResourceGroupNameArg, AAZStrArg
+        args_schema = {}
+        args_schema.cluster_name = AAZStrArg(
+            options=["--cluster-name", "-n"],
+            help="The name of the cluster.",
+            required=True
+        )
+        args_schema.resource_group = AAZResourceGroupNameArg(
+            options=["--resource-group", "-g"],
+            help="Name of resource group.",
+            required=True
+        )
+        args_schema.vm_name = AAZStrArg(
+            options=["--vm-name"],
+            help="The name of the virtual machine.",
+            required=True
+        )
+        return args_schema
+    
+    def __call__(self, cmd, **kwargs):
+        from azure.cli.core.commands.client_factory import get_subscription_id
+        from azure.cli.core.util import send_raw_request
+        import json
+        
+        cluster_name = kwargs.get('cluster_name')
+        resource_group = kwargs.get('resource_group')
+        vm_name = kwargs.get('vm_name')
+        
+        subscription_id = get_subscription_id(cmd.cli_ctx)
+        
+        # Construct the REST API path for deprovision
+        path = f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers/Microsoft.AzureStackHCI/clusters/{cluster_name}/jobs/VmConnectDeprovision"
+        
+        # API version
+        api_version = "2023-12-01-preview"
+        url = f"https://management.azure.com{path}?api-version={api_version}"
+        
+        # Payload for VM Connect deprovision
+        payload = {
+            "properties": {
+                "jobType": "VmConnectDeprovision",
+                "deploymentMode": "Deploy",
+                "vmConnectProvisionJobDetails": [
+                    {
+                        "vmName": vm_name
+                    }
+                ]
+            }
+        }
+        
+        # Make the REST API call
+        try:
+            response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
+            return response.json() if response.content else {"message": f"VM Connect deprovision job initiated successfully for VM: {vm_name}"}
+        except Exception as e:
+            from azure.cli.core.util import CLIError
+            raise CLIError(f"Failed to disable VM Connect for VM '{vm_name}': {str(e)}")

From d939b594760d4ba4b90fc279b87b85f67aba3c53 Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Tue, 29 Jul 2025 22:38:59 -0400
Subject: [PATCH 02/16] Update src/stack-hci/azext_stack_hci/custom.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 src/stack-hci/azext_stack_hci/custom.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/stack-hci/azext_stack_hci/custom.py b/src/stack-hci/azext_stack_hci/custom.py
index c62b891bdbf..047793609d9 100644
--- a/src/stack-hci/azext_stack_hci/custom.py
+++ b/src/stack-hci/azext_stack_hci/custom.py
@@ -85,7 +85,7 @@ def __init__(self, loader=None):
     
     @classmethod
     def _build_arguments_schema(cls, *args, **kwargs):
-        from azure.cli.core.aaz import AAZResourceGroupNameArg, AAZStrArg, AAZDictArg
+        from azure.cli.core.aaz import AAZResourceGroupNameArg, AAZStrArg
         args_schema = {}
         args_schema.cluster_name = AAZStrArg(
             options=["--cluster-name", "-n"],

From 79de363dd14291dc84137325a78db7ace950a2c0 Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Wed, 30 Jul 2025 02:41:53 +0000
Subject: [PATCH 03/16] updates

---
 src/stack-hci/azext_stack_hci/custom.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/src/stack-hci/azext_stack_hci/custom.py b/src/stack-hci/azext_stack_hci/custom.py
index c62b891bdbf..6615d8b5569 100644
--- a/src/stack-hci/azext_stack_hci/custom.py
+++ b/src/stack-hci/azext_stack_hci/custom.py
@@ -119,9 +119,7 @@ def __call__(self, cmd, **kwargs):
         path = f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers/Microsoft.AzureStackHCI/clusters/{cluster_name}/jobs/VmConnectProvision"
         
         # API version
-        api_version = "2023-12-01-preview"
-        url = f"https://management.azure.com{path}?api-version={api_version}"
-        
+        url = f"https://management.azure.com{path}?api-version={API_VERSION}"
         # Default payload with VM name
         payload = {
             "properties": {
@@ -192,9 +190,9 @@ def __call__(self, cmd, **kwargs):
         # Payload for VM Connect deprovision
         payload = {
             "properties": {
-                "jobType": "VmConnectDeprovision",
+                "jobType": "VmConnectRemove",
                 "deploymentMode": "Deploy",
-                "vmConnectProvisionJobDetails": [
+                "vmConnectRemoveJobDetails": [
                     {
                         "vmName": vm_name
                     }

From 71df8981f497e449b3711937f883b2195689319a Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Wed, 30 Jul 2025 02:48:24 +0000
Subject: [PATCH 04/16] fix api version

---
 src/stack-hci/azext_stack_hci/custom.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/stack-hci/azext_stack_hci/custom.py b/src/stack-hci/azext_stack_hci/custom.py
index 8c55328731e..0c1046ebf2c 100644
--- a/src/stack-hci/azext_stack_hci/custom.py
+++ b/src/stack-hci/azext_stack_hci/custom.py
@@ -119,7 +119,8 @@ def __call__(self, cmd, **kwargs):
         path = f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers/Microsoft.AzureStackHCI/clusters/{cluster_name}/jobs/VmConnectProvision"
         
         # API version
-        url = f"https://management.azure.com{path}?api-version={API_VERSION}"
+        api_version = "2023-12-01-preview"
+        url = f"https://management.azure.com{path}?api-version={api_version}"
         # Default payload with VM name
         payload = {
             "properties": {

From 2fa60766135d90b6f1e1766d1c1b319aec894928 Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Wed, 30 Jul 2025 15:43:48 +0000
Subject: [PATCH 05/16] Fix linter errors

---
 src/stack-hci/azext_stack_hci/custom.py | 21 +++++----------------
 1 file changed, 5 insertions(+), 16 deletions(-)

diff --git a/src/stack-hci/azext_stack_hci/custom.py b/src/stack-hci/azext_stack_hci/custom.py
index 0c1046ebf2c..438001bb756 100644
--- a/src/stack-hci/azext_stack_hci/custom.py
+++ b/src/stack-hci/azext_stack_hci/custom.py
@@ -77,16 +77,11 @@ def pre_instance_update(self, instance):
             args.type = 'None'
 
 
-class VmConnectEnable:
-    """Enable VM Connect functionality for Stack HCI cluster."""
-    
-    def __init__(self, loader=None):
-        self.loader = loader
-    
+class VmConnectEnable(_VmConnectEnable):
     @classmethod
     def _build_arguments_schema(cls, *args, **kwargs):
         from azure.cli.core.aaz import AAZResourceGroupNameArg, AAZStrArg
-        args_schema = {}
+        args_schema = super()._build_arguments_schema(*args, **kwargs)
         args_schema.cluster_name = AAZStrArg(
             options=["--cluster-name", "-n"],
             help="The name of the cluster.",
@@ -111,8 +106,7 @@ def __call__(self, cmd, **kwargs):
         
         cluster_name = kwargs.get('cluster_name')
         resource_group = kwargs.get('resource_group')
-        vm_name = kwargs.get('vm_name')
-        
+        vm_name = kwargs.get('vm_name')      
         subscription_id = get_subscription_id(cmd.cli_ctx)
         
         # Construct the REST API path
@@ -143,16 +137,11 @@ def __call__(self, cmd, **kwargs):
             raise CLIError(f"Failed to enable VM Connect for VM '{vm_name}': {str(e)}")
 
 
-class VmConnectDisable:
-    """Disable VM Connect functionality for Stack HCI cluster."""
-    
-    def __init__(self, loader=None):
-        self.loader = loader
-    
+class VmConnectDisable(_VmConnectDisable):
     @classmethod
     def _build_arguments_schema(cls, *args, **kwargs):
         from azure.cli.core.aaz import AAZResourceGroupNameArg, AAZStrArg
-        args_schema = {}
+        args_schema = super()._build_arguments_schema(*args, **kwargs)
         args_schema.cluster_name = AAZStrArg(
             options=["--cluster-name", "-n"],
             help="The name of the cluster.",

From 837963da28abe3acf1d726dc3e878f6e7899569b Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Wed, 30 Jul 2025 16:34:38 +0000
Subject: [PATCH 06/16] removing trailing whitespaces

---
 src/stack-hci/azext_stack_hci/custom.py | 26 ++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/src/stack-hci/azext_stack_hci/custom.py b/src/stack-hci/azext_stack_hci/custom.py
index 438001bb756..dcfd7626418 100644
--- a/src/stack-hci/azext_stack_hci/custom.py
+++ b/src/stack-hci/azext_stack_hci/custom.py
@@ -98,20 +98,20 @@ def _build_arguments_schema(cls, *args, **kwargs):
             required=True
         )
         return args_schema
-    
+
     def __call__(self, cmd, **kwargs):
         from azure.cli.core.commands.client_factory import get_subscription_id
         from azure.cli.core.util import send_raw_request
         import json
-        
+
         cluster_name = kwargs.get('cluster_name')
         resource_group = kwargs.get('resource_group')
-        vm_name = kwargs.get('vm_name')      
+        vm_name = kwargs.get('vm_name')
         subscription_id = get_subscription_id(cmd.cli_ctx)
-        
+
         # Construct the REST API path
         path = f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers/Microsoft.AzureStackHCI/clusters/{cluster_name}/jobs/VmConnectProvision"
-        
+
         # API version
         api_version = "2023-12-01-preview"
         url = f"https://management.azure.com{path}?api-version={api_version}"
@@ -127,7 +127,7 @@ def __call__(self, cmd, **kwargs):
                 ]
             }
         }
-        
+
         # Make the REST API call
         try:
             response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
@@ -158,25 +158,25 @@ def _build_arguments_schema(cls, *args, **kwargs):
             required=True
         )
         return args_schema
-    
+
     def __call__(self, cmd, **kwargs):
         from azure.cli.core.commands.client_factory import get_subscription_id
         from azure.cli.core.util import send_raw_request
         import json
-        
+
         cluster_name = kwargs.get('cluster_name')
         resource_group = kwargs.get('resource_group')
         vm_name = kwargs.get('vm_name')
-        
+
         subscription_id = get_subscription_id(cmd.cli_ctx)
-        
+
         # Construct the REST API path for deprovision
         path = f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers/Microsoft.AzureStackHCI/clusters/{cluster_name}/jobs/VmConnectDeprovision"
-        
+
         # API version
         api_version = "2023-12-01-preview"
         url = f"https://management.azure.com{path}?api-version={api_version}"
-        
+
         # Payload for VM Connect deprovision
         payload = {
             "properties": {
@@ -189,7 +189,7 @@ def __call__(self, cmd, **kwargs):
                 ]
             }
         }
-        
+
         # Make the REST API call
         try:
             response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))

From 12b2485b6d0d5343232b81ba0bf4b0c1a8d84434 Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Wed, 30 Jul 2025 18:01:27 +0000
Subject: [PATCH 07/16] fix line length

---
 src/stack-hci/azext_stack_hci/custom.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/stack-hci/azext_stack_hci/custom.py b/src/stack-hci/azext_stack_hci/custom.py
index dcfd7626418..4d52ff6a1af 100644
--- a/src/stack-hci/azext_stack_hci/custom.py
+++ b/src/stack-hci/azext_stack_hci/custom.py
@@ -110,8 +110,11 @@ def __call__(self, cmd, **kwargs):
         subscription_id = get_subscription_id(cmd.cli_ctx)
 
         # Construct the REST API path
-        path = f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers/Microsoft.AzureStackHCI/clusters/{cluster_name}/jobs/VmConnectProvision"
-
+        path = (
+            f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/"
+            "providers/Microsoft.AzureStackHCI/clusters/"
+            f"{cluster_name}/jobs/VmConnectDeprovision"
+        )
         # API version
         api_version = "2023-12-01-preview"
         url = f"https://management.azure.com{path}?api-version={api_version}"
@@ -171,7 +174,11 @@ def __call__(self, cmd, **kwargs):
         subscription_id = get_subscription_id(cmd.cli_ctx)
 
         # Construct the REST API path for deprovision
-        path = f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers/Microsoft.AzureStackHCI/clusters/{cluster_name}/jobs/VmConnectDeprovision"
+        path = (
+            f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/"
+            "providers/Microsoft.AzureStackHCI/clusters/"
+            f"{cluster_name}/jobs/VmConnectDeprovision"
+        )
 
         # API version
         api_version = "2023-12-01-preview"

From 63aeaefd2e71afa85dbd6a019bebaf457e1f1990 Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Wed, 30 Jul 2025 18:20:08 +0000
Subject: [PATCH 08/16] fix

---
 src/stack-hci/azext_stack_hci/custom.py | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/stack-hci/azext_stack_hci/custom.py b/src/stack-hci/azext_stack_hci/custom.py
index 4d52ff6a1af..d62a351f476 100644
--- a/src/stack-hci/azext_stack_hci/custom.py
+++ b/src/stack-hci/azext_stack_hci/custom.py
@@ -134,7 +134,14 @@ def __call__(self, cmd, **kwargs):
         # Make the REST API call
         try:
             response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
-            return response.json() if response.content else {"message": f"VM Connect provision job initiated successfully for VM: {vm_name}"}
+            if response.content:
+                return response.json()
+            else:
+                return {
+                    "message": (
+                        f"VM Connect provision job initiated successfully for VM: {vm_name}"
+                    )
+                }
         except Exception as e:
             from azure.cli.core.util import CLIError
             raise CLIError(f"Failed to enable VM Connect for VM '{vm_name}': {str(e)}")
@@ -200,7 +207,14 @@ def __call__(self, cmd, **kwargs):
         # Make the REST API call
         try:
             response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
-            return response.json() if response.content else {"message": f"VM Connect deprovision job initiated successfully for VM: {vm_name}"}
+            if response.content:
+                return response.json()
+            else:
+                return {
+                    "message": (
+                        f"VM Connect provision job initiated successfully for VM: {vm_name}"
+                    )
+                }
         except Exception as e:
             from azure.cli.core.util import CLIError
             raise CLIError(f"Failed to disable VM Connect for VM '{vm_name}': {str(e)}")

From 570a2851e4470021eeea9fcfab54d746915ca9dc Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Wed, 30 Jul 2025 22:24:52 +0000
Subject: [PATCH 09/16] fix

---
 src/stack-hci/azext_stack_hci/custom.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/stack-hci/azext_stack_hci/custom.py b/src/stack-hci/azext_stack_hci/custom.py
index d62a351f476..73ab43be68e 100644
--- a/src/stack-hci/azext_stack_hci/custom.py
+++ b/src/stack-hci/azext_stack_hci/custom.py
@@ -113,7 +113,7 @@ def __call__(self, cmd, **kwargs):
         path = (
             f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/"
             "providers/Microsoft.AzureStackHCI/clusters/"
-            f"{cluster_name}/jobs/VmConnectDeprovision"
+            f"{cluster_name}/jobs/VmConnectRemove"
         )
         # API version
         api_version = "2023-12-01-preview"
@@ -180,18 +180,18 @@ def __call__(self, cmd, **kwargs):
 
         subscription_id = get_subscription_id(cmd.cli_ctx)
 
-        # Construct the REST API path for deprovision
+        # Construct the REST API path for Remove
         path = (
             f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/"
             "providers/Microsoft.AzureStackHCI/clusters/"
-            f"{cluster_name}/jobs/VmConnectDeprovision"
+            f"{cluster_name}/jobs/VmConnectRemove"
         )
 
         # API version
         api_version = "2023-12-01-preview"
         url = f"https://management.azure.com{path}?api-version={api_version}"
 
-        # Payload for VM Connect deprovision
+        # Payload for VM Connect Remove
         payload = {
             "properties": {
                 "jobType": "VmConnectRemove",

From 976cf8fd33f679b101fa5ee819c03670ac2a992b Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Wed, 30 Jul 2025 22:25:49 +0000
Subject: [PATCH 10/16] fix

---
 src/stack-hci/azext_stack_hci/custom.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/stack-hci/azext_stack_hci/custom.py b/src/stack-hci/azext_stack_hci/custom.py
index 73ab43be68e..f9837a218d0 100644
--- a/src/stack-hci/azext_stack_hci/custom.py
+++ b/src/stack-hci/azext_stack_hci/custom.py
@@ -212,7 +212,7 @@ def __call__(self, cmd, **kwargs):
             else:
                 return {
                     "message": (
-                        f"VM Connect provision job initiated successfully for VM: {vm_name}"
+                        f"VM Connect remove job initiated successfully for VM: {vm_name}"
                     )
                 }
         except Exception as e:

From ad22a990507899c16b04055d9d3aa731ae8b06fa Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Wed, 30 Jul 2025 22:48:08 +0000
Subject: [PATCH 11/16] fix

---
 src/stack-hci/azext_stack_hci/custom.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/stack-hci/azext_stack_hci/custom.py b/src/stack-hci/azext_stack_hci/custom.py
index f9837a218d0..d92fc01b3eb 100644
--- a/src/stack-hci/azext_stack_hci/custom.py
+++ b/src/stack-hci/azext_stack_hci/custom.py
@@ -77,7 +77,10 @@ def pre_instance_update(self, instance):
             args.type = 'None'
 
 
-class VmConnectEnable(_VmConnectEnable):
+class VmConnectEnable():
+    def __init__(self, loader=None):
+        self.loader = loader
+        
     @classmethod
     def _build_arguments_schema(cls, *args, **kwargs):
         from azure.cli.core.aaz import AAZResourceGroupNameArg, AAZStrArg
@@ -147,7 +150,10 @@ def __call__(self, cmd, **kwargs):
             raise CLIError(f"Failed to enable VM Connect for VM '{vm_name}': {str(e)}")
 
 
-class VmConnectDisable(_VmConnectDisable):
+class VmConnectDisable():
+    def __init__(self, loader=None):
+        self.loader = loader
+
     @classmethod
     def _build_arguments_schema(cls, *args, **kwargs):
         from azure.cli.core.aaz import AAZResourceGroupNameArg, AAZStrArg

From 8b5b2c70dbf5137b7a0ec4562782cd46b5ec2299 Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Thu, 31 Jul 2025 16:18:08 +0000
Subject: [PATCH 12/16] fix

---
 src/stack-hci/azext_stack_hci/custom.py | 24 +++++++++++-------------
 1 file changed, 11 insertions(+), 13 deletions(-)

diff --git a/src/stack-hci/azext_stack_hci/custom.py b/src/stack-hci/azext_stack_hci/custom.py
index d92fc01b3eb..5689a58770d 100644
--- a/src/stack-hci/azext_stack_hci/custom.py
+++ b/src/stack-hci/azext_stack_hci/custom.py
@@ -80,7 +80,7 @@ def pre_instance_update(self, instance):
 class VmConnectEnable():
     def __init__(self, loader=None):
         self.loader = loader
-        
+
     @classmethod
     def _build_arguments_schema(cls, *args, **kwargs):
         from azure.cli.core.aaz import AAZResourceGroupNameArg, AAZStrArg
@@ -139,12 +139,11 @@ def __call__(self, cmd, **kwargs):
             response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
             if response.content:
                 return response.json()
-            else:
-                return {
-                    "message": (
-                        f"VM Connect provision job initiated successfully for VM: {vm_name}"
-                    )
-                }
+            return {
+                "message": (
+                    f"VM Connect provision job initiated successfully for VM: {vm_name}"
+                )
+            }
         except Exception as e:
             from azure.cli.core.util import CLIError
             raise CLIError(f"Failed to enable VM Connect for VM '{vm_name}': {str(e)}")
@@ -215,12 +214,11 @@ def __call__(self, cmd, **kwargs):
             response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
             if response.content:
                 return response.json()
-            else:
-                return {
-                    "message": (
-                        f"VM Connect remove job initiated successfully for VM: {vm_name}"
-                    )
-                }
+            return {
+                "message": (
+                    f"VM Connect remove job initiated successfully for VM: {vm_name}"
+                )
+            }
         except Exception as e:
             from azure.cli.core.util import CLIError
             raise CLIError(f"Failed to disable VM Connect for VM '{vm_name}': {str(e)}")

From 7a62fd5c285998eae9d900fe05134ba950f6cf94 Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Thu, 31 Jul 2025 18:31:26 +0000
Subject: [PATCH 13/16] fixes

---
 src/stack-hci/azext_stack_hci/commands.py |   6 +-
 src/stack-hci/azext_stack_hci/custom.py   | 209 +++++++---------------
 2 files changed, 69 insertions(+), 146 deletions(-)

diff --git a/src/stack-hci/azext_stack_hci/commands.py b/src/stack-hci/azext_stack_hci/commands.py
index fcf39640117..c0c40f10fad 100644
--- a/src/stack-hci/azext_stack_hci/commands.py
+++ b/src/stack-hci/azext_stack_hci/commands.py
@@ -20,7 +20,7 @@ def load_command_table(self, _):  # pylint: disable=unused-argument
         self.command_table['stack-hci cluster identity assign'] = IdentityAssign(loader=self)
         self.command_table['stack-hci cluster identity remove'] = IdentityRemove(loader=self)
 
-    with self.command_group('stack-hci vmconnect'):
+    with self.command_group('stack-hci vm-connect', ...) as g:
         from azext_stack_hci.custom import VmConnectEnable, VmConnectDisable
-        self.command_table['stack-hci vmconnect enable'] = VmConnectEnable(loader=self)
-        self.command_table['stack-hci vmconnect disable'] = VmConnectDisable(loader=self)
+        g.custom_command('enable', 'vm_connect_enable')
+        g.custom_command('disable', 'vm_connect_disable')
diff --git a/src/stack-hci/azext_stack_hci/custom.py b/src/stack-hci/azext_stack_hci/custom.py
index 5689a58770d..41ceb90cd04 100644
--- a/src/stack-hci/azext_stack_hci/custom.py
+++ b/src/stack-hci/azext_stack_hci/custom.py
@@ -77,148 +77,71 @@ def pre_instance_update(self, instance):
             args.type = 'None'
 
 
-class VmConnectEnable():
-    def __init__(self, loader=None):
-        self.loader = loader
-
-    @classmethod
-    def _build_arguments_schema(cls, *args, **kwargs):
-        from azure.cli.core.aaz import AAZResourceGroupNameArg, AAZStrArg
-        args_schema = super()._build_arguments_schema(*args, **kwargs)
-        args_schema.cluster_name = AAZStrArg(
-            options=["--cluster-name", "-n"],
-            help="The name of the cluster.",
-            required=True
-        )
-        args_schema.resource_group = AAZResourceGroupNameArg(
-            options=["--resource-group", "-g"],
-            help="Name of resource group.",
-            required=True
-        )
-        args_schema.vm_name = AAZStrArg(
-            options=["--vm-name"],
-            help="The name of the virtual machine.",
-            required=True
-        )
-        return args_schema
-
-    def __call__(self, cmd, **kwargs):
-        from azure.cli.core.commands.client_factory import get_subscription_id
-        from azure.cli.core.util import send_raw_request
-        import json
-
-        cluster_name = kwargs.get('cluster_name')
-        resource_group = kwargs.get('resource_group')
-        vm_name = kwargs.get('vm_name')
-        subscription_id = get_subscription_id(cmd.cli_ctx)
-
-        # Construct the REST API path
-        path = (
-            f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/"
-            "providers/Microsoft.AzureStackHCI/clusters/"
-            f"{cluster_name}/jobs/VmConnectRemove"
-        )
-        # API version
-        api_version = "2023-12-01-preview"
-        url = f"https://management.azure.com{path}?api-version={api_version}"
-        # Default payload with VM name
-        payload = {
-            "properties": {
-                "jobType": "VmConnectProvision",
-                "deploymentMode": "Deploy",
-                "vmConnectProvisionJobDetails": [
-                    {
-                        "vmName": vm_name
-                    }
-                ]
-            }
+def vm_connect_enable(cmd, cluster_name, resource_group, vm_name):
+    """
+    Enable VM Connect for a virtual machine in an Azure Stack HCI cluster.
+    """
+    from azure.cli.core.commands.client_factory import get_subscription_id
+    from azure.cli.core.util import send_raw_request
+    import json
+
+    subscription_id = get_subscription_id(cmd.cli_ctx)
+    path = (
+        f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/"
+        f"providers/Microsoft.AzureStackHCI/clusters/{cluster_name}/jobs/VmConnectProvision"
+    )
+    api_version = "2023-12-01-preview"
+    url = f"https://management.azure.com{path}?api-version={api_version}"
+    payload = {
+        "properties": {
+            "jobType": "VmConnectProvision",
+            "deploymentMode": "Deploy",
+            "vmConnectProvisionJobDetails": [
+                {"vmName": vm_name}
+            ]
         }
-
-        # Make the REST API call
-        try:
-            response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
-            if response.content:
-                return response.json()
-            return {
-                "message": (
-                    f"VM Connect provision job initiated successfully for VM: {vm_name}"
-                )
-            }
-        except Exception as e:
-            from azure.cli.core.util import CLIError
-            raise CLIError(f"Failed to enable VM Connect for VM '{vm_name}': {str(e)}")
-
-
-class VmConnectDisable():
-    def __init__(self, loader=None):
-        self.loader = loader
-
-    @classmethod
-    def _build_arguments_schema(cls, *args, **kwargs):
-        from azure.cli.core.aaz import AAZResourceGroupNameArg, AAZStrArg
-        args_schema = super()._build_arguments_schema(*args, **kwargs)
-        args_schema.cluster_name = AAZStrArg(
-            options=["--cluster-name", "-n"],
-            help="The name of the cluster.",
-            required=True
-        )
-        args_schema.resource_group = AAZResourceGroupNameArg(
-            options=["--resource-group", "-g"],
-            help="Name of resource group.",
-            required=True
-        )
-        args_schema.vm_name = AAZStrArg(
-            options=["--vm-name"],
-            help="The name of the virtual machine.",
-            required=True
-        )
-        return args_schema
-
-    def __call__(self, cmd, **kwargs):
-        from azure.cli.core.commands.client_factory import get_subscription_id
-        from azure.cli.core.util import send_raw_request
-        import json
-
-        cluster_name = kwargs.get('cluster_name')
-        resource_group = kwargs.get('resource_group')
-        vm_name = kwargs.get('vm_name')
-
-        subscription_id = get_subscription_id(cmd.cli_ctx)
-
-        # Construct the REST API path for Remove
-        path = (
-            f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/"
-            "providers/Microsoft.AzureStackHCI/clusters/"
-            f"{cluster_name}/jobs/VmConnectRemove"
-        )
-
-        # API version
-        api_version = "2023-12-01-preview"
-        url = f"https://management.azure.com{path}?api-version={api_version}"
-
-        # Payload for VM Connect Remove
-        payload = {
-            "properties": {
-                "jobType": "VmConnectRemove",
-                "deploymentMode": "Deploy",
-                "vmConnectRemoveJobDetails": [
-                    {
-                        "vmName": vm_name
-                    }
-                ]
-            }
+    }
+
+    try:
+        response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
+        if response.content:
+            return response.json()
+        return {"message": f"VM Connect provision job initiated successfully for VM: {vm_name}"}
+    except Exception as e:
+        from azure.cli.core.util import CLIError
+        raise CLIError(f"Failed to enable VM Connect for VM '{vm_name}': {str(e)}")
+
+
+def vm_connect_disable(cmd, cluster_name, resource_group, vm_name):
+    """
+    Disable VM Connect for a virtual machine in an Azure Stack HCI cluster.
+    """
+    from azure.cli.core.commands.client_factory import get_subscription_id
+    from azure.cli.core.util import send_raw_request
+    import json
+
+    subscription_id = get_subscription_id(cmd.cli_ctx)
+    path = (
+        f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/"
+        f"providers/Microsoft.AzureStackHCI/clusters/{cluster_name}/jobs/VmConnectRemove"
+    )
+    api_version = "2023-12-01-preview"
+    url = f"https://management.azure.com{path}?api-version={api_version}"
+    payload = {
+        "properties": {
+            "jobType": "VmConnectRemove",
+            "deploymentMode": "Deploy",
+            "vmConnectRemoveJobDetails": [
+                {"vmName": vm_name}
+            ]
         }
-
-        # Make the REST API call
-        try:
-            response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
-            if response.content:
-                return response.json()
-            return {
-                "message": (
-                    f"VM Connect remove job initiated successfully for VM: {vm_name}"
-                )
-            }
-        except Exception as e:
-            from azure.cli.core.util import CLIError
-            raise CLIError(f"Failed to disable VM Connect for VM '{vm_name}': {str(e)}")
+    }
+
+    try:
+        response = send_raw_request(cmd.cli_ctx, "PUT", url, body=json.dumps(payload))
+        if response.content:
+            return response.json()
+        return {"message": f"VM Connect remove job initiated successfully for VM: {vm_name}"}
+    except Exception as e:
+        from azure.cli.core.util import CLIError
+        raise CLIError(f"Failed to disable VM Connect for VM '{vm_name}': {str(e)}")

From de846d4db4fa35fc4a02144be785e7836d18b4c5 Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Thu, 31 Jul 2025 20:30:55 +0000
Subject: [PATCH 14/16] fix

---
 src/stack-hci/azext_stack_hci/commands.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/stack-hci/azext_stack_hci/commands.py b/src/stack-hci/azext_stack_hci/commands.py
index c0c40f10fad..ca11fd3b92e 100644
--- a/src/stack-hci/azext_stack_hci/commands.py
+++ b/src/stack-hci/azext_stack_hci/commands.py
@@ -21,6 +21,5 @@ def load_command_table(self, _):  # pylint: disable=unused-argument
         self.command_table['stack-hci cluster identity remove'] = IdentityRemove(loader=self)
 
     with self.command_group('stack-hci vm-connect', ...) as g:
-        from azext_stack_hci.custom import VmConnectEnable, VmConnectDisable
         g.custom_command('enable', 'vm_connect_enable')
         g.custom_command('disable', 'vm_connect_disable')

From 3f79d75c52536bb5be95b599381ca49a93599bff Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Thu, 31 Jul 2025 20:40:28 +0000
Subject: [PATCH 15/16] fix

---
 src/stack-hci/azext_stack_hci/_help.py    | 21 +++++++++++++++++++++
 src/stack-hci/azext_stack_hci/commands.py |  6 +++---
 src/stack-hci/azext_stack_hci/custom.py   |  4 ++--
 3 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/src/stack-hci/azext_stack_hci/_help.py b/src/stack-hci/azext_stack_hci/_help.py
index 126d5d00714..61cfc95e451 100644
--- a/src/stack-hci/azext_stack_hci/_help.py
+++ b/src/stack-hci/azext_stack_hci/_help.py
@@ -9,3 +9,24 @@
 # pylint: disable=too-many-lines
 
 from knack.help_files import helps  # pylint: disable=unused-import
+
+helps['stack-hci vmconnect'] = """
+type: group
+short-summary: Manage VM Connect for Azure Stack HCI virtual machines.
+"""
+
+helps['stack-hci vmconnect enable'] = """
+type: command
+short-summary: Enable VM Connect for a virtual machine in an Azure Stack HCI cluster.
+examples:
+  - name: Enable VM Connect
+    text: az stack-hci vmconnect enable --cluster-name MyCluster --resource-group MyResourceGroup --vm-name MyVM
+"""
+
+helps['stack-hci vmconnect disable'] = """
+type: command
+short-summary: Disable VM Connect for a virtual machine in an Azure Stack HCI cluster.
+examples:
+  - name: Disable VM Connect
+    text: az stack-hci vmconnect disable --cluster-name MyCluster --resource-group MyResourceGroup --vm-name MyVM
+"""
diff --git a/src/stack-hci/azext_stack_hci/commands.py b/src/stack-hci/azext_stack_hci/commands.py
index ca11fd3b92e..da3b895e7d1 100644
--- a/src/stack-hci/azext_stack_hci/commands.py
+++ b/src/stack-hci/azext_stack_hci/commands.py
@@ -20,6 +20,6 @@ def load_command_table(self, _):  # pylint: disable=unused-argument
         self.command_table['stack-hci cluster identity assign'] = IdentityAssign(loader=self)
         self.command_table['stack-hci cluster identity remove'] = IdentityRemove(loader=self)
 
-    with self.command_group('stack-hci vm-connect', ...) as g:
-        g.custom_command('enable', 'vm_connect_enable')
-        g.custom_command('disable', 'vm_connect_disable')
+    with self.command_group('stack-hci vmconnect', ...) as g:
+        g.custom_command('enable', 'vmconnect_enable')
+        g.custom_command('disable', 'vmconnect_disable')
diff --git a/src/stack-hci/azext_stack_hci/custom.py b/src/stack-hci/azext_stack_hci/custom.py
index 41ceb90cd04..c04126854a7 100644
--- a/src/stack-hci/azext_stack_hci/custom.py
+++ b/src/stack-hci/azext_stack_hci/custom.py
@@ -77,7 +77,7 @@ def pre_instance_update(self, instance):
             args.type = 'None'
 
 
-def vm_connect_enable(cmd, cluster_name, resource_group, vm_name):
+def vmconnect_enable(cmd, cluster_name, resource_group, vm_name):
     """
     Enable VM Connect for a virtual machine in an Azure Stack HCI cluster.
     """
@@ -112,7 +112,7 @@ def vm_connect_enable(cmd, cluster_name, resource_group, vm_name):
         raise CLIError(f"Failed to enable VM Connect for VM '{vm_name}': {str(e)}")
 
 
-def vm_connect_disable(cmd, cluster_name, resource_group, vm_name):
+def vmconnect_disable(cmd, cluster_name, resource_group, vm_name):
     """
     Disable VM Connect for a virtual machine in an Azure Stack HCI cluster.
     """

From 0a4bcd0795980f8c207af63ad49813dc10005be6 Mon Sep 17 00:00:00 2001
From: hvedati <43424696+hvedati@users.noreply.github.com>
Date: Thu, 31 Jul 2025 21:49:24 +0000
Subject: [PATCH 16/16] adding params

---
 src/stack-hci/azext_stack_hci/_params.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/stack-hci/azext_stack_hci/_params.py b/src/stack-hci/azext_stack_hci/_params.py
index cfcec717c9c..0a09eff68fd 100644
--- a/src/stack-hci/azext_stack_hci/_params.py
+++ b/src/stack-hci/azext_stack_hci/_params.py
@@ -10,4 +10,12 @@
 
 
 def load_arguments(self, _):  # pylint: disable=unused-argument
-    pass
+    with self.argument_context('stack-hci vmconnect enable') as c:
+        c.argument('cluster_name', options_list=['--cluster-name', '-n'], help='The name of the cluster.')
+        c.argument('resource_group', options_list=['--resource-group', '-g'], help='Name of resource group.')
+        c.argument('vm_name', options_list=['--vm-name'], help='The name of the virtual machine.')
+
+    with self.argument_context('stack-hci vmconnect disable') as c:
+        c.argument('cluster_name', options_list=['--cluster-name', '-n'], help='The name of the cluster.')
+        c.argument('resource_group', options_list=['--resource-group', '-g'], help='Name of resource group.')
+        c.argument('vm_name', options_list=['--vm-name'], help='The name of the virtual machine.')