From 0cd390fae9b2437e753c156341e66bc367037486 Mon Sep 17 00:00:00 2001
From: Jakub Krajewski <jakubpkrajewski@gmail.com>
Date: Wed, 27 Mar 2024 17:28:27 +0100
Subject: [PATCH] Add device template model with id and devices attached

---
 catalystwan/api/template_api.py                       |  5 +++--
 .../api/templates/device_template/device_template.py  |  5 +++++
 catalystwan/models/configuration/config_migration.py  |  4 ++--
 catalystwan/workflows/config_migration.py             | 11 +++++++----
 4 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/catalystwan/api/template_api.py b/catalystwan/api/template_api.py
index d8a7dd73..8bb080dd 100644
--- a/catalystwan/api/template_api.py
+++ b/catalystwan/api/template_api.py
@@ -16,6 +16,7 @@
 from catalystwan.api.templates.device_template.device_template import (
     DeviceSpecificValue,
     DeviceTemplate,
+    DeviceTemplateWithTracking,
     GeneralTemplate,
 )
 from catalystwan.api.templates.feature_template import FeatureTemplate
@@ -747,7 +748,7 @@ def get_device_templates(self) -> DataSequence[DeviceTemplateInformation]:
         templates = self.session.get(url=endpoint, params=params)
         return templates.dataseq(DeviceTemplateInformation)
 
-    def get_device_template(self, template_id: str) -> DeviceTemplate:
+    def get_device_template(self, template_id: str) -> DeviceTemplateWithTracking:
         endpoint = f"/dataservice/template/device/object/{template_id}"
         response = self.session.get(endpoint)
-        return DeviceTemplate(**response.json())
+        return DeviceTemplateWithTracking(**response.json())
diff --git a/catalystwan/api/templates/device_template/device_template.py b/catalystwan/api/templates/device_template/device_template.py
index 4cc6d95b..d6b6e7b6 100644
--- a/catalystwan/api/templates/device_template/device_template.py
+++ b/catalystwan/api/templates/device_template/device_template.py
@@ -108,5 +108,10 @@ def get(self, name: str, session: ManagerSession) -> DeviceTemplate:
     model_config = ConfigDict(populate_by_name=True, use_enum_values=True)
 
 
+class DeviceTemplateWithTracking(DeviceTemplate):
+    id: str = Field(default="", alias="templateId")
+    devices_attached: int = Field(default=0, alias="devicesAttached")
+
+
 class DeviceSpecificValue(BaseModel):
     property: str
diff --git a/catalystwan/models/configuration/config_migration.py b/catalystwan/models/configuration/config_migration.py
index 54520bf8..13a1c668 100644
--- a/catalystwan/models/configuration/config_migration.py
+++ b/catalystwan/models/configuration/config_migration.py
@@ -7,7 +7,7 @@
 from typing_extensions import Annotated
 
 from catalystwan.api.template_api import FeatureTemplateInformation
-from catalystwan.api.templates.device_template.device_template import DeviceTemplate
+from catalystwan.api.templates.device_template.device_template import DeviceTemplateWithTracking
 from catalystwan.endpoints.configuration_group import ConfigGroupCreationPayload
 from catalystwan.models.configuration.feature_profile.common import FeatureProfileCreationPayload, ProfileType
 from catalystwan.models.configuration.feature_profile.sdwan.other import AnyOtherParcel
@@ -50,7 +50,7 @@ class UX1Templates(BaseModel):
     feature_templates: List[FeatureTemplateInformation] = Field(
         default=[], serialization_alias="featureTemplates", validation_alias="featureTemplates"
     )
-    device_templates: List[DeviceTemplate] = Field(
+    device_templates: List[DeviceTemplateWithTracking] = Field(
         default=[], serialization_alias="deviceTemplates", validation_alias="deviceTemplates"
     )
 
diff --git a/catalystwan/workflows/config_migration.py b/catalystwan/workflows/config_migration.py
index 57a7cf8e..e8eab214 100644
--- a/catalystwan/workflows/config_migration.py
+++ b/catalystwan/workflows/config_migration.py
@@ -210,10 +210,13 @@ def collect_ux1_config(session: ManagerSession, progress: Callable[[str, int, in
     ux1.templates.feature_templates = [t for t in template_api.get_feature_templates()]
     progress("Collecting Feature Templates", 1, 2)
 
-    device_templates_ids = [t.id for t in template_api.get_device_templates()]
-    for i, dtid in enumerate(device_templates_ids):
-        ux1.templates.device_templates.append(template_api.get_device_template(dtid))
-        progress("Collecting Device Templates", i + 1, len(device_templates_ids))
+    device_templates_information = template_api.get_device_templates()
+    for i, device_template_information in enumerate(device_templates_information):
+        device_template = template_api.get_device_template(device_template_information.id)
+        device_template.id = device_template_information.id
+        device_template.devices_attached = device_template_information.devices_attached
+        ux1.templates.device_templates.append(device_template)
+        progress("Collecting Device Templates", i + 1, len(device_templates_information))
 
     return ux1