Skip to content
This repository has been archived by the owner on Apr 26, 2024. It is now read-only.

Commit

Permalink
Move function flatten_general_templates to DeviceTemplate model. (#526)
Browse files Browse the repository at this point in the history
* Move function flatten_general_templates to DeviceTemplate model.

* Rename
  • Loading branch information
jpkrajewski authored Mar 21, 2024
1 parent 60a4a45 commit 43cd647
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 0 deletions.
20 changes: 20 additions & 0 deletions catalystwan/api/templates/device_template/device_template.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,26 @@ class DeviceTemplate(BaseModel):
security_policy_id: str = Field(default="", alias="securityPolicyId")
policy_id: str = Field(default="", alias="policyId")

def get_flattened_general_templates(self) -> List[GeneralTemplate]:
"""
Recursively flattens the general templates by removing the sub-templates
and returning a list of flattened templates.
Returns:
A list of GeneralTemplate objects representing the flattened templates.
"""

def flatten_general_templates(general_templates: List[GeneralTemplate]) -> List[GeneralTemplate]:
result = []
for gt in general_templates:
sub_templates = gt.subTemplates
gt.subTemplates = []
result.append(gt)
result.extend(flatten_general_templates(sub_templates))
return result

return flatten_general_templates(self.general_templates)

def generate_payload(self) -> str:
env = Environment(
loader=FileSystemLoader(self.payload_path.parent),
Expand Down
48 changes: 48 additions & 0 deletions catalystwan/tests/templates/test_device_template.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
import unittest

from catalystwan.api.templates.device_template.device_template import DeviceTemplate, GeneralTemplate


class TestDeviceTemplate(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(
self.device_template.get_flattened_general_templates(),
[
GeneralTemplate(
name="1level",
templateId="1",
templateType="1",
),
GeneralTemplate(
name="2level",
templateId="2",
templateType="2",
),
GeneralTemplate(name="3level", templateId="3", templateType="3"),
],
)

0 comments on commit 43cd647

Please sign in to comment.