feat: Init includes default dataflow profile (#294)

* Expose one parameter --df-profile-instances = 1

azext_edge/edge/commands_edge.py

@@ -124,6 +124,7 @@ def init(
     mq_authn_name: str = "authn",
     mq_broker_config_file: Optional[str] = None,
     mq_insecure: Optional[bool] = None,
+    dataflow_profile_instances: int = 1,
     disable_secret_rotation: Optional[bool] = None,
     rotation_poll_interval: str = "1h",
     csi_driver_version: str = KEYVAULT_ARC_EXTENSION_VERSION,
@@ -224,6 +225,7 @@ def init(
         mq_broker_name=str(mq_broker_name),
         mq_authn_name=str(mq_authn_name),
         mq_insecure=mq_insecure,
+        dataflow_profile_instances=int(dataflow_profile_instances),
         keyvault_resource_id=keyvault_resource_id,
         keyvault_spc_secret_name=str(keyvault_spc_secret_name),
         disable_secret_rotation=disable_secret_rotation,

azext_edge/edge/params.py

@@ -619,6 +619,13 @@ def load_iotops_arguments(self, _):
             help="The path to a custom IoT Operations deployment template. Intended for advanced use cases.",
             deprecate_info=context.deprecate(hide=True),
         )
+        context.argument(
+            "dataflow_profile_instances",
+            type=int,
+            options_list=["--df-profile-instances"],
+            help="The instance count associated with the default dataflow profile.",
+            arg_group="Dataflow Profile",
+        )
 
     with self.argument_context("iot ops delete") as context:
         context.argument(

azext_edge/edge/providers/orchestration/template.py

@@ -519,3 +519,22 @@ def get_current_template_copy(custom_template_path: Optional[str] = None) -> Tem
         moniker=CURRENT_TEMPLATE.moniker,
         content=deepcopy(CURRENT_TEMPLATE.content),
     )
+
+
+def get_basic_dataflow_profile(profile_name: str = "profile", instance_count: int = 1):
+    return {
+        "type": "Microsoft.IoTOperations/instances/dataflowProfiles",
+        "apiVersion": "2024-07-01-preview",
+        "name": f"[format('{{0}}/{{1}}', parameters('instanceName'), '{profile_name}')]",
+        "extendedLocation": {
+            "name": "[resourceId('Microsoft.ExtendedLocation/customLocations', parameters('customLocationName'))]",
+            "type": "CustomLocation",
+        },
+        "properties": {
+            "instanceCount": instance_count,
+        },
+        "dependsOn": [
+            "[resourceId('Microsoft.IoTOperations/instances', parameters('instanceName'))]",
+            "[resourceId('Microsoft.ExtendedLocation/customLocations', parameters('customLocationName'))]",
+        ],
+    }

azext_edge/edge/providers/orchestration/work.py

@@ -7,7 +7,7 @@
 from enum import IntEnum
 from json import dumps
 from time import sleep
-from typing import Dict, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 from uuid import uuid4
 
 from azure.cli.core.azclierror import AzureResponseError, ValidationError
@@ -22,7 +22,12 @@
 
 from ...util import get_timestamp_now_utc
 from ...util.x509 import DEFAULT_EC_ALGO, DEFAULT_VALID_DAYS
-from .template import CURRENT_TEMPLATE, TemplateVer, get_current_template_copy
+from .template import (
+    CURRENT_TEMPLATE,
+    TemplateVer,
+    get_basic_dataflow_profile,
+    get_current_template_copy,
+)
 
 logger = get_logger(__name__)
 
@@ -663,6 +668,11 @@ def build_template(self, work_kpis: dict) -> Tuple[TemplateVer, dict]:
             broker: dict = template.get_resource_defs("Microsoft.IoTOperations/instances/brokers")
             broker["properties"] = mq_broker_config
 
+        # Default dataflow profile
+        deploy_resources: List[dict] = template.content.get("resources", [])
+        df_profile_instances = self._kwargs.get("dataflow_profile_instances", 1)
+        deploy_resources.append(get_basic_dataflow_profile(instance_count=df_profile_instances))
+
         return template, parameters
 
 

azext_edge/tests/edge/init/int/dataflow_helper.py

@@ -0,0 +1,25 @@
+# coding=utf-8
+# ----------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License file in the project root for license information.
+# ----------------------------------------------------------------------------------------------
+
+from typing import List, Optional
+
+from .helper import get_resource_from_partial_id, ResourceKeys
+
+
+def assert_dataflow_profile_args(
+    resource_group: str, init_resources: List[str], dataflow_profile_instances: Optional[int] = None, **_
+):
+    instance_resources = [res for res in init_resources if res.startswith(ResourceKeys.iot_operations.value)]
+    instance_partial_id = instance_resources[0]
+    instance_resources = set(instance_resources)
+
+    # Dataflow Profile
+    expected_profile_partial_id = f"{instance_partial_id}/dataflowProfiles/profile"
+    assert expected_profile_partial_id in instance_resources
+
+    profile = get_resource_from_partial_id(expected_profile_partial_id, resource_group)
+    profile_props = profile["properties"]
+    assert profile_props["instanceCount"] == (1 or dataflow_profile_instances)

azext_edge/tests/edge/init/int/helper.py

@@ -155,6 +155,7 @@ def _assert_instance_show(instance_name: str, resource_group: str, ops_version:
 
     assert show_result["extendedLocation"]["name"].endswith(arg_dict.get("custom_location", "-ops-init-cl"))
     assert show_result["extendedLocation"]["type"] == "CustomLocation"
+
     assert show_result["location"] == arg_dict.get(
         "location", run(f"az group show -n {resource_group}")["location"]
     ).lower()

azext_edge/tests/edge/init/int/test_init_int.py

@@ -4,15 +4,18 @@
 # Licensed under the MIT License. See License file in the project root for license information.
 # ----------------------------------------------------------------------------------------------
 
+from os import mkdir
 from typing import Dict, Union
+
 import pytest
 from knack.log import get_logger
-from os import mkdir
+
 from ....helpers import run
+from .dataflow_helper import assert_dataflow_profile_args
+from .helper import assert_init_result, strip_quotes
 from .mq_helper import assert_broker_args
 from .opcua_helper import assert_simulate_plc_args
 from .orchestrator_helper import assert_orchestrator_args
-from .helper import assert_init_result, strip_quotes
 
 logger = get_logger(__name__)
 
@@ -58,7 +61,7 @@ def init_test_setup(cluster_connection, settings):
 def test_init_scenario(
     init_test_setup, tracked_files
 ):
-    additional_args = init_test_setup["additionalArgs"]
+    additional_args = init_test_setup["additionalArgs"] or ""
     arg_dict = _process_additional_args(additional_args)
 
     if "ca_dir" in arg_dict:
@@ -112,6 +115,7 @@ def test_init_scenario(
         for assertion in [
             assert_simulate_plc_args,
             assert_broker_args,
+            assert_dataflow_profile_args,
             assert_orchestrator_args
         ]:
             assertion(
@@ -137,7 +141,7 @@ def _process_additional_args(additional_args: str) -> Dict[str, Union[str, bool]
     arg_dict = {}
     for arg in additional_args.split("--")[1:]:
         arg = arg.strip().split(" ", maxsplit=1)
-        # --simualte-plc vs --desc "potato cluster"
+        # --simulate-plc vs --desc "potato cluster"
         arg[0] = arg[0].replace("-", "_")
         if len(arg) == 1 or arg[1].lower() == "true":
             arg_dict[arg[0]] = True

azext_edge/tests/edge/init/test_work_unit.py

@@ -31,6 +31,7 @@
     WorkCategoryKey,
     WorkManager,
     WorkStepKey,
+    get_basic_dataflow_profile,
 )
 from azext_edge.edge.util import assemble_nargs_to_dict
 
@@ -76,6 +77,7 @@ def mock_broker_config():
     mq_insecure,
     disable_rsync_rules,
     mq_broker_config_file,
+    dataflow_profile_instances,
     """,
     [
         pytest.param(
@@ -106,6 +108,7 @@ def mock_broker_config():
             None,  # mq_insecure
             None,  # disable_rsync_rules
             None,  # mq_broker_config_file
+            None,  # dataflow_profile_instances
         ),
         pytest.param(
             None,  # instance_name
@@ -135,6 +138,7 @@ def mock_broker_config():
             None,  # mq_insecure
             None,  # disable_rsync_rules
             None,  # mq_broker_config_file
+            None,  # dataflow_profile_instances
         ),
         pytest.param(
             generate_random_string(),  # instance_name
@@ -164,6 +168,7 @@ def mock_broker_config():
             True,  # mq_insecure
             True,  # disable_rsync_rules
             str(MOCK_BROKER_CONFIG_PATH),  # mq_broker_config_file
+            randint(5, 10),  # dataflow_profile_instances
         ),
     ],
 )
@@ -199,6 +204,7 @@ def test_init_to_template_params(
     mq_insecure,
     disable_rsync_rules,
     mq_broker_config_file,
+    dataflow_profile_instances,
 ):
     kwargs = {}
 
@@ -228,6 +234,7 @@ def test_init_to_template_params(
         (mq_insecure, "mq_insecure"),
         (disable_rsync_rules, "disable_rsync_rules"),
         (mq_broker_config_file, "mq_broker_config_file"),
+        (dataflow_profile_instances, "dataflow_profile_instances"),
     ]
 
     for param_tuple in param_tuples:
@@ -344,6 +351,16 @@ def test_init_to_template_params(
     if mq_broker_config_file:
         assert broker["properties"] == mock_broker_config
 
+    dataflow_profiles = template_ver.get_resource_defs(
+        resource_type="Microsoft.IoTOperations/instances/dataflowProfiles", first=False
+    )
+    assert len(dataflow_profiles) == 1
+    dataflow_profile = dataflow_profiles[0]
+    df_profile_kwargs = {}
+    if dataflow_profile_instances:
+        df_profile_kwargs["instance_count"] = dataflow_profile_instances
+    assert dataflow_profile == get_basic_dataflow_profile(**df_profile_kwargs)
+
 
 @pytest.mark.parametrize(
     """

azext_edge/tests/edge/orchestration/test_template_unit.py

@@ -23,7 +23,7 @@ def test_current_template():
 
     assert CURRENT_TEMPLATE.get_component_vers() == CURRENT_TEMPLATE.content["variables"]["VERSIONS"]
 
-    for r_type in ["brokers", "brokers/listeners"]:
+    for r_type in ["brokers", "brokers/listeners", "brokers/authentications"]:
         fqr_type = f"Microsoft.IoTOperations/instances/{r_type}"
         assert CURRENT_TEMPLATE.get_resource_defs(fqr_type)["type"] == fqr_type