From 69580a3ef1c5e72a4fc42a81861021390b5d54fe Mon Sep 17 00:00:00 2001
From: Jakub Krajewski <jakubpkrajewski@gmail.com>
Date: Thu, 28 Mar 2024 10:48:34 +0100
Subject: [PATCH] Remove deprecated function, use class method. Minor change:
 use real uuid in config group from device template

---
 .../test_flatten_general_templates.py         | 49 -------------------
 .../device_templates/__init__.py              | 11 -----
 .../flatten_general_templates.py              | 22 ---------
 catalystwan/workflows/config_migration.py     |  5 +-
 4 files changed, 2 insertions(+), 85 deletions(-)
 delete mode 100644 catalystwan/tests/config_migration/test_flatten_general_templates.py
 delete mode 100644 catalystwan/utils/config_migration/device_templates/__init__.py
 delete mode 100644 catalystwan/utils/config_migration/device_templates/flatten_general_templates.py

diff --git a/catalystwan/tests/config_migration/test_flatten_general_templates.py b/catalystwan/tests/config_migration/test_flatten_general_templates.py
deleted file mode 100644
index 72a49fdb..00000000
--- a/catalystwan/tests/config_migration/test_flatten_general_templates.py
+++ /dev/null
@@ -1,49 +0,0 @@
-import unittest
-
-from catalystwan.api.templates.device_template.device_template import DeviceTemplate, GeneralTemplate
-from catalystwan.utils.config_migration.device_templates import flatten_general_templates
-
-
-class TestFlattenGeneralTemplates(unittest.TestCase):
-    def setUp(self):
-        self.device_template = DeviceTemplate(
-            template_name="DT-example",
-            template_description="DT-example",
-            device_role="None",
-            device_type="None",
-            security_policy_id="None",
-            policy_id="None",
-            generalTemplates=[
-                GeneralTemplate(
-                    name="1level",
-                    templateId="1",
-                    templateType="1",
-                    subTemplates=[
-                        GeneralTemplate(
-                            name="2level",
-                            templateId="2",
-                            templateType="2",
-                            subTemplates=[GeneralTemplate(name="3level", templateId="3", templateType="3")],
-                        )
-                    ],
-                )
-            ],
-        )
-
-    def test_flatten_general_templates(self):
-        self.assertEqual(
-            flatten_general_templates(self.device_template.general_templates),
-            [
-                GeneralTemplate(
-                    name="1level",
-                    templateId="1",
-                    templateType="1",
-                ),
-                GeneralTemplate(
-                    name="2level",
-                    templateId="2",
-                    templateType="2",
-                ),
-                GeneralTemplate(name="3level", templateId="3", templateType="3"),
-            ],
-        )
diff --git a/catalystwan/utils/config_migration/device_templates/__init__.py b/catalystwan/utils/config_migration/device_templates/__init__.py
deleted file mode 100644
index adf137c2..00000000
--- a/catalystwan/utils/config_migration/device_templates/__init__.py
+++ /dev/null
@@ -1,11 +0,0 @@
-from typing import List
-
-from .flatten_general_templates import flatten_general_templates
-
-__all__ = [
-    "flatten_general_templates",
-]
-
-
-def __dir__() -> "List[str]":
-    return list(__all__)
diff --git a/catalystwan/utils/config_migration/device_templates/flatten_general_templates.py b/catalystwan/utils/config_migration/device_templates/flatten_general_templates.py
deleted file mode 100644
index b9fc97af..00000000
--- a/catalystwan/utils/config_migration/device_templates/flatten_general_templates.py
+++ /dev/null
@@ -1,22 +0,0 @@
-from typing import List
-
-from catalystwan.api.templates.device_template.device_template import GeneralTemplate
-
-
-def flatten_general_templates(general_templates: List[GeneralTemplate]) -> List[GeneralTemplate]:
-    """
-    Recursively flattens a list of GeneralTemplate objects.
-
-    Args:
-        general_templates (List[GeneralTemplate]): The list of GeneralTemplate objects to flatten.
-
-    Returns:
-        List[GeneralTemplate]: The flattened list of GeneralTemplate objects.
-    """
-    result = []
-    for gt in general_templates:
-        sub_templates = gt.subTemplates
-        gt.subTemplates = []
-        result.append(gt)
-        result.extend(flatten_general_templates(sub_templates))
-    return result
diff --git a/catalystwan/workflows/config_migration.py b/catalystwan/workflows/config_migration.py
index 789a9830..a9c6aa9f 100644
--- a/catalystwan/workflows/config_migration.py
+++ b/catalystwan/workflows/config_migration.py
@@ -18,7 +18,6 @@
 from catalystwan.utils.config_migration.converters.feature_template import create_parcel_from_template
 from catalystwan.utils.config_migration.converters.policy.policy_lists import convert as convert_policy_list
 from catalystwan.utils.config_migration.creators.config_pusher import UX2ConfigPusher, UX2ConfigRollback
-from catalystwan.utils.config_migration.device_templates import flatten_general_templates
 from catalystwan.utils.config_migration.reverters.config_reverter import UX2ConfigReverter
 
 logger = logging.getLogger(__name__)
@@ -96,7 +95,7 @@ def transform(ux1: UX1Config) -> UX2Config:
     ux2 = UX2Config()
     # Create Feature Profiles and Config Group
     for dt in ux1.templates.device_templates:
-        templates = flatten_general_templates(dt.general_templates)
+        templates = dt.get_flattened_general_templates()
 
         # Create Feature Profiles
         fp_system_uuid = uuid4()
@@ -132,7 +131,7 @@ def transform(ux1: UX1Config) -> UX2Config:
         transformed_cg = TransformedConfigGroup(
             header=TransformHeader(
                 type="config_group",
-                origin=uuid4(),
+                origin=UUID(dt.template_id),
                 subelements=set([fp_system_uuid, fp_other_uuid]),
             ),
             config_group=ConfigGroupCreationPayload(