[load] Adding support for creating and managing Schedule in load test resource using CLI (#8540)

* Adding commands for scheduling

* added list schedule command

* initial tests for scheduling

* Updated tests

* Enhancing test cases for schedule

* Optimizing utils and addressing PR comments

* Adding invalid test cases

* Addressing PR comments

* linting and styling

* test cases and recordings updated

* nit

* Updating version

* nit

* Addressing PR comments and issues after bug bash

* Fixes after running githooks

* removed datetime.UTC as not supported in python 3.9

Author: Himanshu49

src/load/HISTORY.rst

@@ -2,6 +2,9 @@
 
 Release History
 ===============
+1.8.0
+++++++
+* Add commands for creating and managing schedule triggers using CLI. 
 
 1.7.0
 ++++++

src/load/azext_load/__init__.py

@@ -18,6 +18,7 @@ def load_command_table(self, args):
         from azext_load.commands import load_command_table
         from azext_load.data_plane.load_test.commands import load_test_commands
         from azext_load.data_plane.load_test_run.commands import load_test_run_commands
+        from azext_load.data_plane.load_trigger.commands import load_trigger_schedule_commands
         from azure.cli.core.aaz import load_aaz_command_table
         try:
             from . import aaz
@@ -32,18 +33,21 @@ def load_command_table(self, args):
         load_command_table(self, args)
         load_test_commands(self, args)
         load_test_run_commands(self, args)
+        load_trigger_schedule_commands(self, args)
         return self.command_table
 
     def load_arguments(self, command):
         from azext_load._params import load_arguments
         from azext_load.data_plane.params import load_arguments as load_common_arguments
         from azext_load.data_plane.load_test.params import load_arguments as load_test_arguments
         from azext_load.data_plane.load_test_run.params import load_arguments as load_test_run_arguments
+        from azext_load.data_plane.load_trigger.params import load_arguments as load_trigger_schedule_arguments
 
         load_arguments(self, command)
         load_common_arguments(self, command)
         load_test_arguments(self, command)
         load_test_run_arguments(self, command)
+        load_trigger_schedule_arguments(self, command)
 
 
 COMMAND_LOADER_CLS = LoadCommandsLoader

src/load/azext_load/_help.py

@@ -12,7 +12,9 @@
 from azext_load.data_plane.help import helps as data_common_helps
 from azext_load.data_plane.load_test.help import helps as data_load_test_helps
 from azext_load.data_plane.load_test_run.help import helps as data_load_test_run_helps
+from azext_load.data_plane.load_trigger.help import helps as data_load_trigger_helps
 
 helps.update(data_common_helps)
 helps.update(data_load_test_helps)
 helps.update(data_load_test_run_helps)
+helps.update(data_load_trigger_helps)

src/load/azext_load/data_plane/help.py

@@ -84,3 +84,19 @@
 short-summary: Command group to retrieve load test run metrics.
 long-summary: Command group to retrieve load test run metrics with list, get-namespaces, get-definitions, get-dimension.
 """ + _common_params
+
+helps[
+    "load trigger"
+] = """
+type: group
+short-summary: Command group to manage trigger.
+long-summary: Command group to manage triggers. Currently the only supported trigger type is schedule.
+""" + _common_params
+
+helps[
+    "load trigger schedule"
+] = """
+type: group
+short-summary: Command group to manage schedule triggers.
+long-summary: Command group to manage schedule triggers.
+""" + _common_params

src/load/azext_load/data_plane/load_trigger/__init__.py

@@ -0,0 +1,6 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

src/load/azext_load/data_plane/load_trigger/commands.py

@@ -0,0 +1,28 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from azure.cli.core.commands import CliCommandType
+
+admin_custom_sdk = CliCommandType(
+    operations_tmpl="azext_load.data_plane.load_trigger.custom#{}"
+)
+
+
+def load_trigger_schedule_commands(self, _):
+    with self.command_group(
+        "load trigger", custom_command_type=admin_custom_sdk, is_preview=True
+    ):
+        pass
+
+    with self.command_group(
+        "load trigger schedule", custom_command_type=admin_custom_sdk, is_preview=True
+    ) as g:
+        g.custom_command("create", "create_trigger_schedule")
+        g.custom_command("update", "update_trigger_schedule")
+        g.custom_command("delete", "delete_trigger_schedule", confirmation=True)
+        g.custom_show_command("show", "get_trigger_schedule")
+        g.custom_command("pause", "pause_trigger_schedule")
+        g.custom_command("enable", "enable_trigger_schedule")
+        g.custom_command("list", "list_trigger_schedules")

src/load/azext_load/data_plane/load_trigger/custom.py

@@ -0,0 +1,247 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+# pylint: disable=line-too-long
+
+from azure.core.exceptions import ResourceNotFoundError
+from azure.cli.core.azclierror import InvalidArgumentValueError, ValidationError
+from knack.log import get_logger
+from azext_load.data_plane.utils.utils import (
+    get_admin_data_plane_client)
+from azext_load.vendored_sdks.loadtesting.models import (_models as models, _enums as enums)
+from azext_load.data_plane.load_trigger import utils
+
+logger = get_logger(__name__)
+
+
+def create_trigger_schedule(
+    cmd,
+    load_test_resource,
+    resource_group_name=None,
+    trigger_id=None,
+    description=None,
+    display_name=None,
+    trigger_start_date_time=None,
+    recurrence_type=None,
+    recurrence_interval=None,
+    recurrence_index=None,
+    recurrence_cron_expression=None,
+    recurrence_dates_in_month=None,
+    recurrence_week_days=None,
+    end_after_occurrence=None,
+    end_after_date_time=None,
+    test_ids=None,
+):
+    client = get_admin_data_plane_client(cmd, load_test_resource, resource_group_name)
+    logger.info("Creating schedule trigger.")
+    try:
+        client.get_trigger(trigger_id)
+        msg = "Trigger schedule with id: {} already exists.".format(trigger_id)
+        logger.error(msg)
+        raise InvalidArgumentValueError(msg)
+    except ResourceNotFoundError:
+        pass
+    recurrence_end_body = utils.get_recurrence_end_body(
+        end_after_occurrence,
+        end_after_date_time,
+    )
+    logger.debug("Recurrence end object: %s", recurrence_end_body)
+    recurrence_body = utils.get_recurrence_body(
+        recurrence_type,
+        recurrence_interval,
+        recurrence_index,
+        recurrence_cron_expression,
+        recurrence_dates_in_month,
+        recurrence_week_days,
+        recurrence_end_body,
+    )
+    logger.debug("Recurrence object: %s", recurrence_body)
+    trigger_body = models.ScheduleTestsTrigger(
+        test_ids=test_ids,
+        recurrence=recurrence_body,
+        start_date_time=trigger_start_date_time,
+        state=enums.TriggerState.ACTIVE,
+        display_name=display_name,
+        description=description,
+    )
+    logger.debug("Trigger schedule body: %s", trigger_body)
+    try:
+        response = client.create_or_update_trigger(trigger_id=trigger_id, body=trigger_body)
+        logger.debug("Created trigger schedule: %s", response)
+        logger.info("Creating trigger schedule completed")
+        return response.as_dict()
+    except Exception:
+        logger.error("Error occurred while creating schedule trigger.")
+        raise
+
+
+def update_trigger_schedule(
+    cmd,
+    load_test_resource,
+    trigger_id,
+    resource_group_name=None,
+    description=None,
+    display_name=None,
+    trigger_start_date_time=None,
+    recurrence_type=None,
+    recurrence_interval=None,
+    recurrence_index=None,
+    recurrence_cron_expression=None,
+    recurrence_dates_in_month=None,
+    recurrence_week_days=None,
+    end_after_occurrence=None,
+    end_after_date_time=None,
+    test_ids=None,
+):
+    client = get_admin_data_plane_client(cmd, load_test_resource, resource_group_name)
+    logger.info("Updating schedule trigger with id: %s", trigger_id)
+    existing_trigger_schedule: models.ScheduleTestsTrigger = None
+    try:
+        existing_trigger_schedule = client.get_trigger(trigger_id)
+    except ResourceNotFoundError:
+        msg = "Schedule trigger with id: {} does not exists.".format(trigger_id)
+        logger.debug(msg)
+        raise InvalidArgumentValueError(msg)
+    logger.debug("Existing schedule trigger: %s", existing_trigger_schedule)
+    recurrence_end_body = utils.get_recurrence_end_body(
+        end_after_occurrence,
+        end_after_date_time,
+        existing_trigger_schedule.recurrence.recurrence_end if existing_trigger_schedule.recurrence else None
+    )
+    logger.debug("Recurrence end object: %s", recurrence_end_body)
+    recurrence_body = utils.get_recurrence_body_for_update(
+        recurrence_type,
+        recurrence_interval,
+        recurrence_index,
+        recurrence_cron_expression,
+        recurrence_dates_in_month,
+        recurrence_week_days,
+        recurrence_end_body,
+        existing_trigger_schedule.recurrence
+    )
+    logger.debug("Recurrence object: %s", recurrence_body)
+    new_trigger_body = models.ScheduleTestsTrigger(
+        test_ids=test_ids,
+        recurrence=recurrence_body,
+        start_date_time=trigger_start_date_time,
+        display_name=display_name,
+        description=description,
+    )
+    new_trigger_body.state = existing_trigger_schedule.state
+    logger.debug("Schedule trigger body to be sent for update: %s", new_trigger_body)
+    try:
+        response = client.create_or_update_trigger(trigger_id=trigger_id, body=new_trigger_body)
+        logger.debug("Updated schedule trigger: %s", response)
+        logger.info("Updating schedule trigger completed")
+        return response.as_dict()
+    except Exception:
+        logger.error("Error occurred while updating schedule trigger.")
+        raise
+
+
+def delete_trigger_schedule(
+    cmd,
+    load_test_resource,
+    trigger_id,
+    resource_group_name=None,
+):
+    logger.info(
+        "Deleting schedule trigger with id: %s", trigger_id
+    )
+    client = get_admin_data_plane_client(cmd, load_test_resource, resource_group_name)
+    client.delete_trigger(trigger_id)
+    logger.info("Deleting schedule trigger completed.")
+
+
+def get_trigger_schedule(
+    cmd,
+    load_test_resource,
+    trigger_id,
+    resource_group_name=None,
+):
+    logger.info(
+        "Getting schedule trigger with id: %s", trigger_id
+    )
+    client = get_admin_data_plane_client(cmd, load_test_resource, resource_group_name)
+    response = client.get_trigger(trigger_id)
+    logger.debug("Fetched schedule trigger: %s", response)
+    return response.as_dict()
+
+
+def pause_trigger_schedule(
+    cmd,
+    load_test_resource,
+    trigger_id,
+    resource_group_name=None,
+):
+    logger.info(
+        "Pausing schedule trigger with id: %s", trigger_id
+    )
+    client = get_admin_data_plane_client(cmd, load_test_resource, resource_group_name)
+    existing_trigger_schedule: models.ScheduleTestsTrigger = None
+    try:
+        existing_trigger_schedule = client.get_trigger(trigger_id)
+    except ResourceNotFoundError:
+        msg = "Schedule trigger with id: {} does not exists.".format(trigger_id)
+        logger.debug(msg)
+        raise InvalidArgumentValueError(msg)
+    logger.debug("Existing schedule trigger object: %s", existing_trigger_schedule)
+    if existing_trigger_schedule.state == enums.TriggerState.ACTIVE:
+        existing_trigger_schedule.state = enums.TriggerState.PAUSED
+        response = client.create_or_update_trigger(trigger_id=trigger_id, body=existing_trigger_schedule)
+        logger.debug("Paused schedule trigger: %s", response)
+        return response.as_dict()
+    if existing_trigger_schedule.state == enums.TriggerState.COMPLETED:
+        msg = "Schedule trigger with id: {} is already completed. A completed schedule cannot be paused.".format(trigger_id)
+        logger.error(msg)
+        raise ValidationError(msg)
+    logger.warning("Schedule trigger is not active. It is in %s state. Enable the schedule before performing pause action.", existing_trigger_schedule.state.value)
+
+
+def enable_trigger_schedule(
+    cmd,
+    load_test_resource,
+    trigger_id,
+    resource_group_name=None,
+):
+    logger.info(
+        "Enabling schedule trigger with id: %s", trigger_id
+    )
+    client = get_admin_data_plane_client(cmd, load_test_resource, resource_group_name)
+    existing_trigger_schedule: models.ScheduleTestsTrigger = None
+    try:
+        existing_trigger_schedule = client.get_trigger(trigger_id)
+    except ResourceNotFoundError:
+        msg = "Schedule trigger with id: {} does not exists.".format(trigger_id)
+        logger.debug(msg)
+        raise InvalidArgumentValueError(msg)
+    logger.debug("Existing trigger object: %s", existing_trigger_schedule)
+    if existing_trigger_schedule.state != enums.TriggerState.COMPLETED:
+        existing_trigger_schedule.state = enums.TriggerState.ACTIVE
+        response = client.create_or_update_trigger(trigger_id=trigger_id, body=existing_trigger_schedule)
+        logger.debug("Enabled schedule trigger: %s", response)
+        return response.as_dict()
+    msg = "Schedule trigger with id: {} is already completed. A completed schedule cannot be enabled.".format(trigger_id)
+    logger.debug(msg)
+    raise ValidationError(msg)
+
+
+def list_trigger_schedules(
+    cmd,
+    load_test_resource,
+    resource_group_name=None,
+    trigger_states=None,
+    test_ids=None,
+):
+    logger.info("Listing schedule triggers.")
+    client = get_admin_data_plane_client(cmd, load_test_resource, resource_group_name)
+    if trigger_states:
+        trigger_states = ",".join(trigger_states)
+    if test_ids:
+        test_ids = ",".join(test_ids)
+    logger.info("Schedule trigger states: %s", trigger_states)
+    response_list = client.list_trigger(test_ids=test_ids, states=trigger_states)
+    logger.debug("Fetched list of schedule triggers: %s", response_list)
+    return [response.as_dict() for response in response_list]