Add files via upload

catalystwan/models/monitoring/security_policy.py

@@ -1,31 +1,16 @@
 # Copyright 2024 Cisco Systems, Inc. and its affiliates
 
 from pydantic import BaseModel, Field
-
-
-class DeviceListResponse(BaseModel):
-    data: dict = Field(..., alias="data")
-
-    def get_amp_down(self) -> list:
-        return self.data["amp_down"]
-
-    def get_amp_up(self) -> list:
-        return self.data["amp_up"]
-
-    def get_ips_down(self) -> list:
-        return self.data["ips_down"]
-
-    def get_ips_up(self) -> list:
-        return self.data["ips_up"]
-
-    def get_urlf_down(self) -> list:
-        return self.data["urlf_down"]
-
-    def get_urlf_up(self) -> list:
-        return self.data["urlf_up"]
-
-    def get_zbfw_down(self) -> list:
-        return self.data["zbfw_down"]
-
-    def get_zbfw_up(self) -> list:
-        return self.data["zbfw_up"]
+from typing import List
+
+
+class SecurityPolicyDeviceList(BaseModel):
+    #data: dict = Field(..., alias="data")
+    amp_down: List[str] = Field(..., alias="amp_down")
+    amp_up: List[str] = Field(..., alias="amp_up")
+    ips_down: List[str] = Field(..., alias="ips_down")
+    ips_up: List[str] = Field(..., alias="ips_up")
+    urlf_down: List[str] = Field(..., alias="urlf_down")
+    urlf_up: List[str] = Field(..., alias="urlf_up")
+    zbfw_down: List[str] = Field(..., serialization_alias="zbfw_down", validation_alias="zbfw_down")
+    zbfw_up: List[str] = Field(..., serialization_alias="zbfw_up", validation_alias="zbfw_up")

catalystwan/tests/test_monitoring_security_policy.py

@@ -11,27 +11,15 @@
 class TestMonitoringSecurityPolicy(unittest.TestCase):
     def setUp(self) -> None:
         self.mock_device_list_data = {
-            "data": {
-                "amp_down": ["2.2.2.2"],
-                "amp_up": ["3.3.3.3"],
-                "ips_down": ["2.2.2.2"],
-                "ips_up": ["1.1.1.1"],
-                "urlf_down": ["3.3.3.3"],
-                "urlf_up": ["1.1.1.1", "2.2.2.2"],
-                "zbfw_down": [],
-                "zbfw_up": ["1.1.1.1", "2.2.2.2", "3.3.3.3"]
-            }
+            "amp_down": ["2.2.2.2"],
+            "amp_up": ["3.3.3.3"],
+            "ips_down": ["2.2.2.2"],
+            "ips_up": ["1.1.1.1"],
+            "urlf_down": ["3.3.3.3"],
+            "urlf_up": ["1.1.1.1", "2.2.2.2"],
+            "zbfw_down": [],
+            "zbfw_up": ["1.1.1.1", "2.2.2.2", "3.3.3.3"]
         }
-        self.mock_device_list_data = SecurityPolicyDeviceList(
-            amp_down=["2.2.2.2"],
-            amp_up=["3.3.3.3"],
-            ips_down=["2.2.2.2"],
-            ips_up=["1.1.1.1"],
-            urlf_down=["3.3.3.3"],
-            urlf_up=["1.1.1.1", "2.2.2.2"],
-            zbfw_down=[],
-            zbfw_up=["1.1.1.1", "2.2.2.2", "3.3.3.3"],
-        )
 
     @patch("catalystwan.session.ManagerSession")
     def test_get_security_policy_device_list(self, mock_session):
@@ -40,7 +28,7 @@ def test_get_security_policy_device_list(self, mock_session):
         mock_response = MagicMock()
 
         # Set up the mock response's dataobj method to return a ServerInfoResponse
-        mock_response.dataobj.return_value = DataSequence(SecurityPolicyDeviceList, [self.mock_device_list_data])
+        mock_response.dataobj.return_value = SecurityPolicyDeviceList(**self.mock_device_list_data)
 
         # Mock the request method of the ManagerSession to return the mock response
         mock_session_instance = mock_session.return_value
@@ -52,23 +40,23 @@ def test_get_security_policy_device_list(self, mock_session):
         response = mon_sec_pol_api.get_device_list()
 
         # Assert
-        self.assertIsInstance(response, DataSequence)
-        self.assertEqual(response[0].amp_down, self.mock_device_list_data.amp_down)
-        self.assertEqual(response[0].amp_down, ["2.2.2.2"])
-        self.assertEqual(response[0].amp_up, self.mock_device_list_data.amp_up)
-        self.assertEqual(response[0].amp_up, ["3.3.3.3"])
-        self.assertEqual(response[0].ips_down, self.mock_device_list_data.ips_down)
-        self.assertEqual(response[0].ips_down, ["2.2.2.2"])
-        self.assertEqual(response[0].ips_up, self.mock_device_list_data.ips_up)
-        self.assertEqual(response[0].ips_up, ["1.1.1.1"])
-        self.assertEqual(response[0].urlf_down, self.mock_device_list_data.urlf_down)
-        self.assertEqual(response[0].urlf_down, ["3.3.3.3"])
-        self.assertEqual(response[0].urlf_up, self.mock_device_list_data.urlf_up)
-        self.assertEqual(response[0].urlf_up, ["1.1.1.1", "2.2.2.2"])
-        self.assertEqual(response[0].zbfw_down, self.mock_device_list_data.zbfw_down)
-        self.assertEqual(response[0].zbfw_down, [])
-        self.assertEqual(response[0].zbfw_up, self.mock_device_list_data.zbfw_up)
-        self.assertEqual(response[0].zbfw_up, ["1.1.1.1", "2.2.2.2", "3.3.3.3"])
+        self.assertIsInstance(response, SecurityPolicyDeviceList)
+        self.assertEqual(response.amp_down, self.mock_device_list_data["amp_down"])
+        self.assertEqual(response.amp_down, ["2.2.2.2"])
+        self.assertEqual(response.amp_up, self.mock_device_list_data["amp_up"])
+        self.assertEqual(response.amp_up, ["3.3.3.3"])
+        self.assertEqual(response.ips_down, self.mock_device_list_data["ips_down"])
+        self.assertEqual(response.ips_down, ["2.2.2.2"])
+        self.assertEqual(response.ips_up, self.mock_device_list_data["ips_up"])
+        self.assertEqual(response.ips_up, ["1.1.1.1"])
+        self.assertEqual(response.urlf_down, self.mock_device_list_data["urlf_down"])
+        self.assertEqual(response.urlf_down, ["3.3.3.3"])
+        self.assertEqual(response.urlf_up, self.mock_device_list_data["urlf_up"])
+        self.assertEqual(response.urlf_up, ["1.1.1.1", "2.2.2.2"])
+        self.assertEqual(response.zbfw_down, self.mock_device_list_data["zbfw_down"])
+        self.assertEqual(response.zbfw_down, [])
+        self.assertEqual(response.zbfw_up, self.mock_device_list_data["zbfw_up"])
+        self.assertEqual(response.zbfw_up, ["1.1.1.1", "2.2.2.2", "3.3.3.3"])
 
         # Ensure the request method was called
         mock_session_instance.request.assert_called_once()