From 96ef2e3b77d6b485d868d79d0f6b1e85939637d5 Mon Sep 17 00:00:00 2001
From: Szymon Basan <sbasan@cisco.com>
Date: Thu, 21 Nov 2024 12:59:16 +0100
Subject: [PATCH] fix: data policy nextHopIpv6 action set entry

---
 catalystwan/models/policy/definition/traffic_data.py | 10 +++++++---
 catalystwan/models/policy/policy_definition.py       |  6 ++++++
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/catalystwan/models/policy/definition/traffic_data.py b/catalystwan/models/policy/definition/traffic_data.py
index 500a8c7d..d1a7b91c 100644
--- a/catalystwan/models/policy/definition/traffic_data.py
+++ b/catalystwan/models/policy/definition/traffic_data.py
@@ -1,6 +1,6 @@
 # Copyright 2023 Cisco Systems, Inc. and its affiliates
 
-from ipaddress import IPv4Address, IPv4Network, IPv6Network
+from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network
 from typing import List, Literal, Optional, Set, Tuple, Union, overload
 from uuid import UUID
 
@@ -49,6 +49,7 @@
     Match,
     NATAction,
     NextHopActionEntry,
+    NextHopIpv6ActionEntry,
     NextHopLooseEntry,
     PacketLengthEntry,
     PLPEntry,
@@ -274,8 +275,11 @@ def associate_nat_action(
         self._insert_action(nat_action)
 
     @accept_action
-    def associate_next_hop_action(self, next_hop: IPv4Address, loose: bool = False) -> None:
-        self._insert_action_in_set(NextHopActionEntry(value=next_hop))
+    def associate_next_hop_action(self, next_hop: Union[IPv4Address, IPv6Address], loose: bool = False) -> None:
+        if isinstance(next_hop, IPv6Address):
+            self._insert_action_in_set(NextHopIpv6ActionEntry(value=next_hop))
+        else:
+            self._insert_action_in_set(NextHopActionEntry(value=next_hop))
         self._insert_action_in_set(NextHopLooseEntry(value=loose))
 
     @accept_action
diff --git a/catalystwan/models/policy/policy_definition.py b/catalystwan/models/policy/policy_definition.py
index 02802294..bcd9d701 100644
--- a/catalystwan/models/policy/policy_definition.py
+++ b/catalystwan/models/policy/policy_definition.py
@@ -492,6 +492,11 @@ class NextHopActionEntry(BaseModel):
     value: Union[IPv4Address, IPv6Address]
 
 
+class NextHopIpv6ActionEntry(BaseModel):
+    field: Literal["nextHopIpv6"] = "nextHopIpv6"
+    value: IPv6Address
+
+
 class NextHopMatchEntry(BaseModel):
     field: Literal["nextHop"] = "nextHop"
     ref: UUID
@@ -1205,6 +1210,7 @@ class CloudSaaSAction(BaseModel):
         MetricEntry,
         MetricTypeEntry,
         NextHopActionEntry,
+        NextHopIpv6ActionEntry,
         NextHopLooseEntry,
         OMPTagEntry,
         OriginatorEntry,