Add new command budget to consumption (#6024)

* Add new command budget to consumption

* Fix build errors on consumption

* Fix test bugs on consumption

* Addressed comments of argument type, wrapper, and comments on consumption

* Fixed build errors on indent

* Fixed import sequence error

* Fixed bugs on argument option list

* Removed enum type validation

src/command_modules/azure-cli-consumption/HISTORY.rst

@@ -2,6 +2,10 @@
 
 Release History
 ===============
+0.3.1
++++++
+* Added new commands for budget API.
+
 0.3.0
 +++++
 * Added commands `marketplace`.

src/command_modules/azure-cli-consumption/azure/cli/command_modules/consumption/_client_factory.py

@@ -28,3 +28,7 @@ def pricesheet_mgmt_client_factory(cli_ctx, kwargs):
 
 def marketplace_mgmt_client_factory(cli_ctx, kwargs):
     return cf_consumption(cli_ctx, **kwargs).marketplaces
+
+
+def budget_mgmt_client_factory(cli_ctx, kwargs):
+    return cf_consumption(cli_ctx, **kwargs).budgets

src/command_modules/azure-cli-consumption/azure/cli/command_modules/consumption/_help.py

@@ -63,3 +63,28 @@
     type: command
     short-summary: List the marketplace for an Azure subscription within a billing period.
 """
+
+helps['consumption budget'] = """
+    type: group
+    short-summary: Manage budgets for an Azure subscription.
+"""
+
+helps['consumption budget list'] = """
+    type: command
+    short-summary: List budgets for an Azure subscription.
+"""
+
+helps['consumption budget show'] = """
+    type: command
+    short-summary: Show budget for an Azure subscription.
+"""
+
+helps['consumption budget create'] = """
+    type: command
+    short-summary: Create a budget for an Azure subscription.
+"""
+
+helps['consumption budget delete'] = """
+    type: command
+    short-summary: Delete a budget for an Azure subscription.
+"""

src/command_modules/azure-cli-consumption/azure/cli/command_modules/consumption/_params.py

@@ -4,33 +4,47 @@
 # --------------------------------------------------------------------------------------------
 
 # pylint: disable=line-too-long
-from ._validators import get_datetime_type
+# pylint: disable=too-many-statements
+from azure.cli.core.commands.parameters import get_enum_type
+from ._validators import (datetime_type,
+                          decimal_type)
 
 
 def load_arguments(self, _):
     with self.argument_context('consumption usage') as c:
         c.argument('top', options_list=['--top', '-t'], type=int, help='Maximum number of items to return. Value range: 1-1000.')
         c.argument('include_additional_properties', options_list=['--include-additional-properties', '-a'], action='store_true', help='Include additional properties in the usages.')
         c.argument('include_meter_details', options_list=['--include-meter-details', '-m'], action='store_true', help='Include meter details in the usages.')
-        c.argument('start_date', options_list=['--start-date', '-s'], type=get_datetime_type(), help='Start date (YYYY-MM-DD in UTC). If specified, also requires --end-date.')
-        c.argument('end_date', options_list=['--end-date', '-e'], type=get_datetime_type(), help='End date (YYYY-MM-DD in UTC). If specified, also requires --start-date.')
+        c.argument('start_date', options_list=['--start-date', '-s'], type=datetime_type, help='Start date (YYYY-MM-DD in UTC). If specified, also requires --end-date.')
+        c.argument('end_date', options_list=['--end-date', '-e'], type=datetime_type, help='End date (YYYY-MM-DD in UTC). If specified, also requires --start-date.')
         c.argument('billing_period_name', options_list=['--billing-period-name', '-p'], help='Name of the billing period to get the usage details that associate with.')
 
     with self.argument_context('consumption reservation') as rs:
-        rs.argument('reservation_order_id', options_list='--reservation-order-id', help='Reservation order id.')
-        rs.argument('start_date', options_list=['--start-date', '-s'], type=get_datetime_type(), help='Start date (YYYY-MM-DD in UTC). Only needed for daily grain and if specified, also requires --end-date.')
-        rs.argument('end_date', options_list=['--end-date', '-e'], type=get_datetime_type(), help='End date (YYYY-MM-DD in UTC). Only needed for daily grain and if specified, also requires --start-date.')
-        rs.argument('reservation_id', options_list='--reservation-id', help='Reservation id.')
+        rs.argument('reservation_order_id', help='Reservation order id.')
+        rs.argument('start_date', options_list=['--start-date', '-s'], type=datetime_type, help='Start date (YYYY-MM-DD in UTC). Only needed for daily grain and if specified, also requires --end-date.')
+        rs.argument('end_date', options_list=['--end-date', '-e'], type=datetime_type, help='End date (YYYY-MM-DD in UTC). Only needed for daily grain and if specified, also requires --start-date.')
+        rs.argument('reservation_id', help='Reservation id.')
 
     with self.argument_context('consumption reservation summary list') as rs:
-        rs.argument('grain', options_list='--grain', type=str, help='Reservation summary grain. Possible values are daily or monthly.')
+        rs.argument('grain', help='Reservation summary grain. Possible values are daily or monthly.')
 
     with self.argument_context('consumption pricesheet show') as cps:
-        cps.argument('include_meter_details', options_list='--include-meter-details', action='store_true', help='Include meter details in the price sheet.')
+        cps.argument('include_meter_details', action='store_true', help='Include meter details in the price sheet.')
         cps.argument('billing_period_name', options_list=['--billing-period-name', '-p'], help='Name of the billing period to get the price sheet.')
 
     with self.argument_context('consumption marketplace list') as cmp:
         cmp.argument('billing_period_name', options_list=['--billing-period-name', '-p'], help='Name of the billing period to get the marketplace.')
         cmp.argument('top', options_list=['--top', '-t'], type=int, help='Maximum number of items to return. Value range: 1-1000.')
-        cmp.argument('start_date', options_list=['--start-date', '-s'], type=get_datetime_type(), help='Start date (YYYY-MM-DD in UTC). If specified, also requires --end-date.')
-        cmp.argument('end_date', options_list=['--end-date', '-e'], type=get_datetime_type(), help='End date (YYYY-MM-DD in UTC). If specified, also requires --start-date.')
+        cmp.argument('start_date', options_list=['--start-date', '-s'], type=datetime_type, help='Start date (YYYY-MM-DD in UTC). If specified, also requires --end-date.')
+        cmp.argument('end_date', options_list=['--end-date', '-e'], type=datetime_type, help='End date (YYYY-MM-DD in UTC). If specified, also requires --start-date.')
+
+    with self.argument_context('consumption budget') as cb:
+        cb.argument('budget_name', help='Name of a budget.')
+        cb.argument('category', arg_type=get_enum_type(['cost', 'usage']), help='Category of the budget can be cost or usage.')
+        cb.argument('amount', type=decimal_type, help='Amount of a budget.')
+        cb.argument('time_grain', arg_type=get_enum_type(['monthly', 'quarterly', 'annually']), help='Time grain of the budget can be monthly, quarterly, or annually.')
+        cb.argument('start_date', options_list=['--start-date', '-s'], type=datetime_type, help='Start date (YYYY-MM-DD in UTC) of time period of a budget.')
+        cb.argument('end_date', options_list=['--end-date', '-e'], type=datetime_type, help='End date (YYYY-MM-DD in UTC) of time period of a budget.')
+        cb.argument('resource_groups', options_list='--resource-group-filter', nargs='+', help='Space-separated list of resource groups to filter on.')
+        cb.argument('resources', options_list='--resource-filter', nargs='+', help='Space-separated list of resource instances to filter on.')
+        cb.argument('meters', options_list='--meter-filter', nargs='+', help='Space-separated list of meters to filter on. Required if category is usage.')

src/command_modules/azure-cli-consumption/azure/cli/command_modules/consumption/_transformers.py

@@ -72,3 +72,22 @@ def marketplace_list_output(result):
 
 def transform_marketplace_list_output(result):
     return [marketplace_list_output(item) for item in result]
+
+
+def budget_output(result):
+    result.amount = str(result.amount)
+    if result.current_spend:
+        result.current_spend.amount = str(result.current_spend.amount)
+    return result
+
+
+def transform_budget_list_output(result):
+    return [budget_output(item) for item in result]
+
+
+def transform_budget_show_output(result):
+    return budget_output(result)
+
+
+def transform_budget_create_update_output(result):
+    return budget_output(result)

src/command_modules/azure-cli-consumption/azure/cli/command_modules/consumption/_validators.py

@@ -4,24 +4,27 @@
 # --------------------------------------------------------------------------------------------
 
 from datetime import datetime
+from decimal import Decimal
 from azure.cli.core.util import CLIError
 
 
-def get_datetime_type():
+def datetime_type(string):
     """ Validates UTC datetime. Examples of accepted forms:
     2017-12-31T01:11:59Z,2017-12-31T01:11Z or 2017-12-31T01Z or 2017-12-31 """
-    def datetime_type(string):
-        """ Validates UTC datetime. Examples of accepted forms:
-        2017-12-31T01:11:59Z,2017-12-31T01:11Z or 2017-12-31T01Z or 2017-12-31 """
-        accepted_date_formats = ['%Y-%m-%dT%H:%M:%SZ', '%Y-%m-%dT%H:%MZ',
-                                 '%Y-%m-%dT%HZ', '%Y-%m-%d']
-        for form in accepted_date_formats:
-            try:
-                return datetime.strptime(string, form)
-            except ValueError:
-                continue
-        raise ValueError("Input '{}' not valid. Valid example: 2017-02-11T23:59:59Z".format(string))
-    return datetime_type
+    accepted_date_formats = ['%Y-%m-%dT%H:%M:%SZ', '%Y-%m-%dT%H:%MZ', '%Y-%m-%dT%HZ', '%Y-%m-%d']
+    for form in accepted_date_formats:
+        try:
+            return datetime.strptime(string, form)
+        except ValueError:
+            continue
+    raise ValueError("Input '{}' not valid. Valid example: 2017-02-11T23:59:59Z".format(string))
+
+
+def decimal_type(string):
+    try:
+        return Decimal(string)
+    except ValueError:
+        raise ValueError("the value passed cannot be converted to decimal")
 
 
 def validate_both_start_end_dates(namespace):
@@ -37,3 +40,8 @@ def validate_reservation_summary(namespace):
         raise CLIError("usage error: --grain  can be either daily or monthly.")
     if data_grain == 'daily' and (not namespace.start_date or not namespace.end_date):
         raise CLIError("usage error: Both --start-date and --end-date need to be supplied for daily grain.")
+
+
+def validate_budget_parameters(namespace):
+    if namespace.amount < 0:
+        raise CLIError("usage error: --amount must be greater than 0")

src/command_modules/azure-cli-consumption/azure/cli/command_modules/consumption/commands.py

@@ -8,15 +8,20 @@
                             transform_reservation_summary_list_output,
                             transform_reservation_detail_list_output,
                             transform_pricesheet_show_output,
-                            transform_marketplace_list_output)
+                            transform_marketplace_list_output,
+                            transform_budget_list_output,
+                            transform_budget_show_output,
+                            transform_budget_create_update_output)
 from ._client_factory import (usage_details_mgmt_client_factory,
                               reservation_summary_mgmt_client_factory,
                               reservation_detail_mgmt_client_factory,
                               pricesheet_mgmt_client_factory,
-                              marketplace_mgmt_client_factory)
+                              marketplace_mgmt_client_factory,
+                              budget_mgmt_client_factory)
 from ._exception_handler import consumption_exception_handler
 from ._validators import (validate_both_start_end_dates,
-                          validate_reservation_summary)
+                          validate_reservation_summary,
+                          validate_budget_parameters)
 
 
 def load_command_table(self, _):
@@ -39,3 +44,12 @@ def load_command_table(self, _):
     with self.command_group('consumption marketplace') as m:
         m.custom_command('list', 'cli_consumption_list_marketplace', transform=transform_marketplace_list_output,
                          exception_handler=consumption_exception_handler, validator=validate_both_start_end_dates, client_factory=marketplace_mgmt_client_factory)
+
+    with self.command_group('consumption budget', exception_handler=consumption_exception_handler, client_factory=budget_mgmt_client_factory) as p:
+        p.custom_command('list', 'cli_consumption_list_budgets', transform=transform_budget_list_output)
+
+        p.custom_command('show', 'cli_consumption_show_budget', transform=transform_budget_show_output)
+
+        p.custom_command('create', 'cli_consumption_create_budget', transform=transform_budget_create_update_output, validator=validate_budget_parameters)
+
+        p.custom_command('delete', 'cli_consumption_delete_budget')

src/command_modules/azure-cli-consumption/azure/cli/command_modules/consumption/custom.py

@@ -89,3 +89,30 @@ def cli_consumption_list_marketplace(client, billing_period_name=None, start_dat
     elif not billing_period_name and top:
         return list(client.list(filter=filter_expression, top=top).advance_page())
     return client.list(filter=filter_expression)
+
+
+def cli_consumption_list_budgets(client, resource_group_name=None):
+    if resource_group_name:
+        return client.list_by_resource_group_name(resource_group_name)
+    return client.list()
+
+
+def cli_consumption_show_budget(client, budget_name, resource_group_name=None):
+    if resource_group_name:
+        return client.get_by_resource_group_name(resource_group_name, budget_name)
+    return client.get(budget_name)
+
+
+def cli_consumption_create_budget(client, budget_name, category, amount, time_grain, start_date, end_date, resource_groups=None, resources=None, meters=None, resource_group_name=None):
+    time_period = client.models.BudgetTimePeriod(start_date, end_date)
+    filters = client.models.Filters(resource_groups=resource_groups, resources=resources, meters=meters)
+    parameters = client.models.Budget(category=category, amount=amount, time_grain=time_grain, time_period=time_period, filters=filters, notifications=None)
+    if resource_group_name:
+        return client.create_or_update(resource_group_name, budget_name, parameters)
+    return client.create_or_update(budget_name, parameters)
+
+
+def cli_consumption_delete_budget(client, budget_name, resource_group_name=None):
+    if resource_group_name:
+        return client.delete(resource_group_name, budget_name)
+    return client.delete(budget_name)

src/command_modules/azure-cli-consumption/azure/cli/command_modules/consumption/tests/latest/recordings/test_consumption_budget_create.yaml

@@ -0,0 +1,34 @@
+interactions:
+- request:
+    body: '{"properties": {"category": "cost", "amount": 100.0, "timeGrain": "monthly",
+      "timePeriod": {"startDate": "2018-02-01T00:00:00.000Z", "endDate": "2018-10-01T00:00:00.000Z"}}}'
+    headers:
+      Accept: [application/json]
+      Accept-Encoding: ['gzip, deflate']
+      CommandName: [consumption budget create]
+      Connection: [keep-alive]
+      Content-Length: ['173']
+      Content-Type: [application/json; charset=utf-8]
+      User-Agent: [python/3.6.4 (Windows-10-10.0.16299-SP0) requests/2.18.4 msrest/0.4.26
+          msrest_azure/0.4.21 azure-mgmt-consumption/2.0.0 Azure-SDK-For-Python AZURECLI/2.0.28]
+      accept-language: [en-US]
+    method: PUT
+    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Consumption/budgets/costbudget?api-version=2018-01-31
+  response:
+    body: {string: '{"id":"subscriptions/0f88eb23-845d-48b9-b363-efd011b05586/providers/Microsoft.Consumption/budgets/costbudget","name":"costbudget","type":"Microsoft.Consumption/budgets","eTag":"\"1d3ab6019183d09\"","properties":{"timePeriod":{"startDate":"2018-02-01T00:00:00Z","endDate":"2018-10-01T00:00:00Z"},"timeGrain":"Monthly","amount":100.0,"currentSpend":null,"category":"Cost","notifications":{},"filters":{"resourceGroups":[],"resources":[],"meters":[]}}}'}
+    headers:
+      cache-control: [no-cache]
+      content-length: ['449']
+      content-type: [application/json; charset=utf-8]
+      date: ['Wed, 21 Feb 2018 22:05:32 GMT']
+      expires: ['-1']
+      location: ['https://consumption.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Consumption/budgets/costbudget?api-version=2018-01-31']
+      pragma: [no-cache]
+      server: [Microsoft-IIS/8.5]
+      session-id: [5215920e-8606-48f6-b900-5b2a5d59794a]
+      strict-transport-security: [max-age=31536000; includeSubDomains]
+      x-content-type-options: [nosniff]
+      x-ms-ratelimit-remaining-subscription-writes: ['1199']
+      x-powered-by: [ASP.NET]
+    status: {code: 201, message: Created}
+version: 1

src/command_modules/azure-cli-consumption/azure/cli/command_modules/consumption/tests/latest/recordings/test_consumption_budget_delete.yaml

@@ -0,0 +1,31 @@
+interactions:
+- request:
+    body: null
+    headers:
+      Accept: [application/json]
+      Accept-Encoding: ['gzip, deflate']
+      CommandName: [consumption budget delete]
+      Connection: [keep-alive]
+      Content-Length: ['0']
+      Content-Type: [application/json; charset=utf-8]
+      User-Agent: [python/3.6.4 (Windows-10-10.0.16299-SP0) requests/2.18.4 msrest/0.4.26
+          msrest_azure/0.4.21 azure-mgmt-consumption/2.0.0 Azure-SDK-For-Python AZURECLI/2.0.28]
+      accept-language: [en-US]
+    method: DELETE
+    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Consumption/budgets/costbudget?api-version=2018-01-31
+  response:
+    body: {string: ''}
+    headers:
+      cache-control: [no-cache]
+      content-length: ['0']
+      date: ['Thu, 22 Feb 2018 00:11:03 GMT']
+      expires: ['-1']
+      pragma: [no-cache]
+      server: [Microsoft-IIS/8.5]
+      session-id: [2fb4175b-6646-45a5-8595-b324cd07818a]
+      strict-transport-security: [max-age=31536000; includeSubDomains]
+      x-content-type-options: [nosniff]
+      x-ms-ratelimit-remaining-subscription-writes: ['1199']
+      x-powered-by: [ASP.NET]
+    status: {code: 200, message: OK}
+version: 1

src/command_modules/azure-cli-consumption/azure/cli/command_modules/consumption/tests/latest/recordings/test_consumption_budget_usage_create.yaml

@@ -0,0 +1,35 @@
+interactions:
+- request:
+    body: '{"properties": {"category": "usage", "amount": 20.0, "timeGrain": "annually",
+      "timePeriod": {"startDate": "2018-02-01T00:00:00.000Z", "endDate": "2018-10-01T00:00:00.000Z"},
+      "filters": {"meters": ["0dfadad2-6e4f-4078-85e1-90c230d4d482"]}}}'
+    headers:
+      Accept: [application/json]
+      Accept-Encoding: ['gzip, deflate']
+      CommandName: [consumption budget create]
+      Connection: [keep-alive]
+      Content-Length: ['239']
+      Content-Type: [application/json; charset=utf-8]
+      User-Agent: [python/3.6.4 (Windows-10-10.0.16299-SP0) requests/2.18.4 msrest/0.4.26
+          msrest_azure/0.4.21 azure-mgmt-consumption/2.0.0 Azure-SDK-For-Python AZURECLI/2.0.28]
+      accept-language: [en-US]
+    method: PUT
+    uri: https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Consumption/budgets/usagetypebudget1?api-version=2018-01-31
+  response:
+    body: {string: '{"id":"subscriptions/0f88eb23-845d-48b9-b363-efd011b05586/providers/Microsoft.Consumption/budgets/usagetypebudget1","name":"usagetypebudget1","type":"Microsoft.Consumption/budgets","eTag":"\"1d3ab875141ad02\"","properties":{"timePeriod":{"startDate":"2018-02-01T00:00:00Z","endDate":"2018-10-01T00:00:00Z"},"timeGrain":"Annually","amount":20.0,"currentSpend":null,"category":"Usage","notifications":{},"filters":{"resourceGroups":[],"resources":[],"meters":["0dfadad2-6e4f-4078-85e1-90c230d4d482"]}}}'}
+    headers:
+      cache-control: [no-cache]
+      content-length: ['500']
+      content-type: [application/json; charset=utf-8]
+      date: ['Thu, 22 Feb 2018 02:46:20 GMT']
+      expires: ['-1']
+      location: ['https://consumption.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/providers/Microsoft.Consumption/budgets/usagetypebudget1?api-version=2018-01-31']
+      pragma: [no-cache]
+      server: [Microsoft-IIS/8.5]
+      session-id: [3268b230-5476-48b4-9e87-07f73dee71d7]
+      strict-transport-security: [max-age=31536000; includeSubDomains]
+      x-content-type-options: [nosniff]
+      x-ms-ratelimit-remaining-subscription-writes: ['1199']
+      x-powered-by: [ASP.NET]
+    status: {code: 201, message: Created}
+version: 1