diff --git a/src/webpubsub/azext_webpubsub/_params.py b/src/webpubsub/azext_webpubsub/_params.py
index 4469c7be66d..d87cfb3ea3a 100644
--- a/src/webpubsub/azext_webpubsub/_params.py
+++ b/src/webpubsub/azext_webpubsub/_params.py
@@ -32,11 +32,11 @@ def load_arguments(self, _):
         c.argument('webpubsub_name', webpubsub_name_type, options_list=['--name', '-n'])
 
     with self.argument_context('webpubsub create') as c:
-        c.argument('sku', help='The sku name of the webpubsub service. Allowed values: Free_F1, Standard_S1')
+        c.argument('sku', help='The sku name of the webpubsub service. Allowed values: Free_F1, Standard_S1, Premium_P1')
         c.argument('unit_count', help='The number of webpubsub service unit count', type=int)
 
     with self.argument_context('webpubsub update') as c:
-        c.argument('sku', help='The sku name of the webpubsub service. Allowed values: Free_F1, Standard_S1')
+        c.argument('sku', help='The sku name of the webpubsub service. Allowed values: Free_F1, Standard_S1, Premium_P1')
         c.argument('unit_count', help='The number of webpubsub service unit count', type=int)
 
     with self.argument_context('webpubsub key regenerate') as c:
diff --git a/src/webpubsub/azext_webpubsub/custom.py b/src/webpubsub/azext_webpubsub/custom.py
index 35dd9d00ad8..de6ba2765d1 100644
--- a/src/webpubsub/azext_webpubsub/custom.py
+++ b/src/webpubsub/azext_webpubsub/custom.py
@@ -53,6 +53,7 @@ def webpubsub_set(client, webpubsub_name, resource_group_name, parameters):
 
 def update_webpubsub(instance, tags=None, sku=None, unit_count=None):
     sku = sku if sku else instance.sku.name
+    unit_count = unit_count if unit_count else instance.sku.capacity
     instance.sku = ResourceSku(name=sku, capacity=unit_count)
 
     if tags is not None: