Add premium tier in help and fix a bug (#5149)

Azure · Jul 26, 2022 · af9f9f1 · af9f9f1
1 parent b0ebdda
commit af9f9f1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git 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
@@ -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: