From 20d74a60113b2ecca6c4986a1523547e89a1ad9b Mon Sep 17 00:00:00 2001
From: cicharka <93913624+cicharka@users.noreply.github.com>
Date: Mon, 26 Feb 2024 11:32:39 +0100
Subject: [PATCH] Enhance validation for SubTaskData (#491)

---
 .../configuration_dashboard_status.py         |  6 +-
 catalystwan/tests/test_task_status_api.py     | 86 +++++++++++++++++++
 2 files changed, 89 insertions(+), 3 deletions(-)

diff --git a/catalystwan/endpoints/configuration_dashboard_status.py b/catalystwan/endpoints/configuration_dashboard_status.py
index 14f0e0828..787438b2d 100644
--- a/catalystwan/endpoints/configuration_dashboard_status.py
+++ b/catalystwan/endpoints/configuration_dashboard_status.py
@@ -1,5 +1,5 @@
 # mypy: disable-error-code="empty-body"
-from typing import List, Optional
+from typing import Dict, List, Optional, Union
 
 from pydantic.v1 import BaseModel, Field
 
@@ -13,7 +13,7 @@ class SubTaskData(BaseModel):
     action: Optional[str]
     activity: List[str]
     current_activity: Optional[str] = Field(alias="currentActivity")
-    action_config: Optional[str] = Field(alias="actionConfig")
+    action_config: Optional[Union[str, Dict]] = Field(alias="actionConfig")
     order: Optional[int]
     uuid: Optional[str]
     hostname: Optional[str] = Field(alias="host-name")
@@ -50,7 +50,7 @@ class Validation(BaseModel):
     rid: Optional[int] = Field(alias="@rid")
     status_id: str = Field(alias="statusId")
     process_id: Optional[str] = Field(alias="processId")
-    action_config: Optional[str] = Field(alias="actionConfig")
+    action_config: Optional[Union[str, Dict]] = Field(alias="actionConfig")
     current_activity: Optional[str] = Field(alias="currentActivity")
     action: Optional[str] = Field(alias="action")
     start_time: Optional[int] = Field(alias="startTime")
diff --git a/catalystwan/tests/test_task_status_api.py b/catalystwan/tests/test_task_status_api.py
index 4dbc3bb2f..5f0d72f70 100644
--- a/catalystwan/tests/test_task_status_api.py
+++ b/catalystwan/tests/test_task_status_api.py
@@ -226,6 +226,80 @@ def setUp(self, mock_session):
             "isCancelEnabled": True,
             "isParallelExecutionEnabled": True,
         }
+        self.response_with_action_config_as_dict = {
+            "data": [
+                {
+                    "local-system-ip": "local_ip",
+                    "statusType": "reboot",
+                    "activity": [],
+                    "system-ip": "system_ip",
+                    "site-id": "siteid",
+                    "uuid": "dev-uuid",
+                    "@rid": 1211,
+                    "personality": "vedge",
+                    "processId": "processid",
+                    "actionConfig": {
+                        "devices": {
+                            "deviceIP": "",
+                            "deviceId": "",
+                            "version": "",
+                            "order": "",
+                            "isNutellaMigration": False,
+                        }
+                    },
+                    "device-type": "vedge",
+                    "action": "reboot",
+                    "startTime": 1685440088317,
+                    "reachability": "reachable",
+                    "order": 0,
+                    "vmanageIP": "vmanage_ip",
+                    "host-name": "vm1",
+                    "version": "vmanage-version",
+                    "deviceID": "deviceid",
+                    "statusId": "success",
+                    "currentActivity": "Done - Reboot",
+                    "deviceModel": "vedge-cloud",
+                    "validity": "valid",
+                    "requestStatus": "received",
+                    "status": "Success",
+                }
+            ],
+            "validation": {
+                "statusType": "reboot",
+                "activity": [],
+                "vmanageIP": "vmanage-ip",
+                "system-ip": "Validation",
+                "deviceID": "Validation",
+                "uuid": "Validation",
+                "@rid": 747,
+                "statusId": "validation_success",
+                "processId": "reboot-9fc30834-cc46-47c5-83c4-0b837cf84f1a",
+                "actionConfig": {
+                    "devices": {"deviceIP": "", "deviceId": "", "version": "", "order": "", "isNutellaMigration": False}
+                },
+                "currentActivity": "Done - Validation",
+                "action": "reboot",
+                "startTime": 1685440057748,
+                "requestStatus": "received",
+                "status": "Validation success",
+                "order": 0,
+            },
+            "summary": {
+                "action": "reboot",
+                "name": "Reboot",
+                "detailsURL": "/dataservice/device/action/status",
+                "startTime": "1685440057829",
+                "endTime": "1685440179295",
+                "userSessionUserName": "admin",
+                "userSessionIP": "10.0.1.1",
+                "tenantName": "DefaultTenant",
+                "total": 1,
+                "status": "done",
+                "count": {"Success": 1},
+            },
+            "isCancelEnabled": True,
+            "isParallelExecutionEnabled": True,
+        }
 
     @patch.object(Task, "_Task__check_validation_status")
     @patch.object(ConfigurationDashboardStatus, "find_status")
@@ -256,6 +330,18 @@ def test_wait_for_completed_empty_data(self, mock_task_response, mock_validation
         # Assert
         self.assertEqual(answer, True)
 
+    @patch.object(Task, "_Task__check_validation_status")
+    @patch.object(ConfigurationDashboardStatus, "find_status")
+    def test_wait_for_completed_with_action_config_as_dict(self, mock_task_response, mock_validation):
+        # Arrange
+        mock_task_response.return_value = TaskData.parse_obj(self.response_with_action_config_as_dict)
+
+        # Act
+        answer = self.task.wait_for_completed(interval_seconds=1).result
+
+        # Assert
+        self.assertEqual(answer, True)
+
     @patch.object(Task, "_Task__check_validation_status")
     @patch.object(ConfigurationDashboardStatus, "find_status")
     def test_wait_for_completed_no_data(self, mock_task_response, mock_validation):