Introduce optional model validation in handle_response

FRINXio · Mar 8, 2024 · 6b4cb4e · 6b4cb4e
1 parent 6c9c07e
commit 6b4cb4e
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/uniconfig/python/frinx_worker/uniconfig/__init__.py b/uniconfig/python/frinx_worker/uniconfig/__init__.py
@@ -1,7 +1,7 @@
 import dataclasses
 import http.client
 import json
-from typing import Any
+from typing import Any, cast
 from typing import Optional
 
 from frinx.common.conductor_enums import TaskResultStatus
@@ -22,15 +22,18 @@ class UniconfigResultDetails(BaseModel):
     task_status: TaskResultStatus
 
 
-def handle_response(response: Response) -> UniconfigResultDetails:
+def handle_response(response: Response, model: Optional[type[BaseModel]] = None) -> UniconfigResultDetails:
     uniconfig_result = UniconfigResultDetails(
         task_status=TaskResultStatus.COMPLETED if response.ok else TaskResultStatus.FAILED,
         logs=f'{response.request.method} request to {response.url} returned with status code {response.status_code}.'
     )
-
     if response.status_code != http.client.NO_CONTENT:
         try:
-            uniconfig_result.output = response.json()
+            response_json = response.json()
+            if model is not None:
+                uniconfig_result.output = model.parse_obj(response_json).dict()
+            else:
+                uniconfig_result.output = response_json
         except json.JSONDecodeError:
             uniconfig_result.logs += 'ERROR: JSON decoding failed - unparsable response content.'
     return uniconfig_result

diff --git a/uniconfig/python/frinx_worker/uniconfig/device_discovery.py b/uniconfig/python/frinx_worker/uniconfig/device_discovery.py
@@ -1,12 +1,12 @@
 from ipaddress import IPv4Address
 from ipaddress import IPv6Address
-from typing import cast
 
 import pydantic
 import requests
 from frinx.common.frinx_rest import UNICONFIG_HEADERS
 from frinx.common.frinx_rest import UNICONFIG_REQUEST_PARAMS
 from frinx.common.frinx_rest import UNICONFIG_URL_BASE
+from frinx.common.type_aliases import DictAny
 from frinx.common.type_aliases import ListStr
 from frinx.common.worker.service import ServiceWorkersImpl
 from frinx.common.worker.service import WorkerImpl
@@ -87,7 +87,7 @@ def validate_udp(cls, udp_port: str) -> list[UdpPortItem] | None:
                     return None
 
         class WorkerOutput(TaskOutput):
-            output: OperationsDiscoverPostResponse
+            output: DictAny
 
         def execute(self, worker_input: WorkerInput) -> TaskResult[WorkerOutput]:
             if Discover.request is None:
@@ -114,11 +114,11 @@ def execute(self, worker_input: WorkerInput) -> TaskResult[WorkerOutput]:
                 ),
             )
 
-            uniconfig_result = handle_response(response)
+            uniconfig_result = handle_response(response, OperationsDiscoverPostResponse)
 
             return TaskResult(
                 status=uniconfig_result.task_status,
                 logs=uniconfig_result.logs,
                 output=self.WorkerOutput(
-                    output=cast(OperationsDiscoverPostResponse,uniconfig_result.output),
+                    output=uniconfig_result.output,
                 ))