use conditionals and reformat some code in modify iptrunk interface workflow

gso/workflows/iptrunk/modify_trunk_interface.py

@@ -2,14 +2,14 @@
 
 import ipaddress
 import json
-from uuid import uuid4
+from uuid import UUID, uuid4
 
 from orchestrator.forms import FormPage, ReadOnlyField
 from orchestrator.forms.validators import UniqueConstrainedList
 from orchestrator.targets import Target
 from orchestrator.types import FormGenerator, State, UUIDstr
 from orchestrator.utils.json import json_dumps
-from orchestrator.workflow import StepList, done, init, step, workflow
+from orchestrator.workflow import StepList, conditional, done, init, step, workflow
 from orchestrator.workflows.steps import resync, store_process_subscription, unsync
 from orchestrator.workflows.utils import wrap_modify_initial_input_form
 from pydantic import validator
@@ -17,6 +17,7 @@ from pydantic_forms.validators import Label
 
 from gso.products.product_blocks.iptrunk import (
     IptrunkInterfaceBlock,
+    IptrunkSideBlock,
     IptrunkType,
     PhyPortCapacity,
 )
@@ -160,22 +161,23 @@ def modify_iptrunk_subscription(
 ) -> State:
     """Modify the subscription in the service database, reflecting the changes to the newly selected interfaces."""
     # Prepare the list of removed AE members
-    previous_ae_members = {}
-    removed_ae_members = {}
-    for side_index in range(2):
-        previous_ae_members[side_index] = [
+    previous_ae_members = [
+        [
             {
                 "interface_name": member.interface_name,
                 "interface_description": member.interface_description,
             }
-            for member in subscription.iptrunk.iptrunk_sides[side_index].iptrunk_side_ae_members
+            for member in side.iptrunk_side_ae_members
         ]
+        for side in subscription.iptrunk.iptrunk_sides
+    ]
+    removed_ae_members = []
+
     for side_index in range(2):
         previous_members = previous_ae_members[side_index]
         current_members = side_a_ae_members if side_index == 0 else side_b_ae_members
-        removed_ae_members[side_index] = [
-            ae_member for ae_member in previous_members if ae_member not in current_members
-        ]
+        removed_ae_members.append([ae_member for ae_member in previous_members if ae_member not in current_members])
+
     subscription.iptrunk.geant_s_sid = geant_s_sid
     subscription.iptrunk.iptrunk_description = iptrunk_description
     subscription.iptrunk.iptrunk_type = iptrunk_type
@@ -267,67 +269,99 @@ def provision_ip_trunk_iface_real(
     return {"subscription": subscription}
 
 
-@step("Update interfaces in Netbox. Reserving interfaces.")
-def update_interfaces_in_netbox(subscription: Iptrunk, removed_ae_members: dict, previous_ae_members: dict) -> State:
-    """Update Netbox such that it contains the new interfaces."""
+def _netbox_update_interfaces(
+    subscription_id: UUID,
+    side_block: IptrunkSideBlock,
+    removed_ae_members: list[dict],
+    previous_ae_members: list[dict],
+) -> None:
     nbclient = NetboxClient()
-    for index, side in enumerate(subscription.iptrunk.iptrunk_sides):
-        if get_router_vendor(side.iptrunk_side_node.owner_subscription_id) == RouterVendor.NOKIA:
-            lag_interface = side.iptrunk_side_ae_iface
-            router_name = side.iptrunk_side_node.router_fqdn
-            # Free removed interfaces
-            for member in removed_ae_members[str(index)]:
-                nbclient.free_interface(router_name, member["interface_name"])
-            # Attach physical interfaces to :term:`LAG`
-            # Update interface description to subscription ID
-            # Reserve interfaces
-            for interface in side.iptrunk_side_ae_members:
-                if any(
-                    ae_member.get("interface_name") == interface.interface_name
-                    for ae_member in previous_ae_members[str(index)]
-                ):
-                    continue
-                nbclient.attach_interface_to_lag(
-                    device_name=side.iptrunk_side_node.router_fqdn,
-                    lag_name=lag_interface,
-                    iface_name=interface.interface_name,
-                    description=str(subscription.subscription_id),
-                )
-                nbclient.reserve_interface(
-                    device_name=side.iptrunk_side_node.router_fqdn,
-                    iface_name=interface.interface_name,
-                )
-    return {
-        "subscription": subscription,
-    }
+
+    # Free removed interfaces
+    for member in removed_ae_members:
+        nbclient.free_interface(side_block.iptrunk_side_node.router_fqdn, member["interface_name"])
+
+    # Attach physical interfaces to :term:`LAG`
+    # Update interface description to subscription ID
+    # Reserve interfaces
+    for interface in side_block.iptrunk_side_ae_members:
+        if any(ae_member["interface_name"] == interface.interface_name for ae_member in previous_ae_members):
+            continue
+        if side_block.iptrunk_side_ae_iface:
+            nbclient.attach_interface_to_lag(
+                device_name=side_block.iptrunk_side_node.router_fqdn,
+                lag_name=side_block.iptrunk_side_ae_iface,
+                iface_name=interface.interface_name,
+                description=str(subscription_id),
+            )
+
+        nbclient.reserve_interface(
+            device_name=side_block.iptrunk_side_node.router_fqdn,
+            iface_name=interface.interface_name,
+        )
+
+
+@step("Netbox: Reserve side A interfaces")
+def netbox_update_interfaces_side_a(
+    subscription: Iptrunk, removed_ae_members: list[list[dict]], previous_ae_members: list[list[dict]]
+) -> None:
+    """Update Netbox such that it contains the new interfaces on side A."""
+    _netbox_update_interfaces(
+        subscription.subscription_id,
+        subscription.iptrunk.iptrunk_sides[0],
+        removed_ae_members[0],
+        previous_ae_members[0],
+    )
+
+
+@step("Netbox: Reserve side B interfaces")
+def netbox_update_interfaces_side_b(
+    subscription: Iptrunk, removed_ae_members: list[list[dict]], previous_ae_members: list[list[dict]]
+) -> None:
+    """Update Netbox such that it contains the new interfaces on side B."""
+    _netbox_update_interfaces(
+        subscription.subscription_id,
+        subscription.iptrunk.iptrunk_sides[1],
+        removed_ae_members[1],
+        previous_ae_members[1],
+    )
 
 
-@step("Allocate interfaces in Netbox")
-def allocate_interfaces_in_netbox(subscription: Iptrunk, previous_ae_members: dict) -> State:
-    """Allocate the :term:`LAG` interfaces in NetBox.
+def _netbox_allocate_interfaces(side_block: IptrunkSideBlock, previous_ae_members: list[dict]) -> None:
+    nbclient = NetboxClient()
 
-    Attach the :term:`LAG` interfaces to the physical interfaces detach old ones from the :term:`LAG`.
+    for interface in side_block.iptrunk_side_ae_members:
+        if any(ae_member["interface_name"] == interface.interface_name for ae_member in previous_ae_members):
+            continue
+        nbclient.allocate_interface(
+            device_name=side_block.iptrunk_side_node.router_fqdn,
+            iface_name=interface.interface_name,
+        )
+
+    # detach the old interfaces from lag
+    if side_block.iptrunk_side_ae_iface:
+        nbclient.detach_interfaces_from_lag(
+            device_name=side_block.iptrunk_side_node.router_fqdn,
+            lag_name=side_block.iptrunk_side_ae_iface,
+        )
+
+
+@step("Netbox: Allocate side A interfaces")
+def allocate_interfaces_in_netbox_side_a(subscription: Iptrunk, previous_ae_members: list[list[dict]]) -> None:
+    """Allocate the :term:`LAG` interfaces on side A in Netbox.
+
+    Attach the :term:`LAG` interface to the physical interface detach old one from the :term:`LAG`.
     """
-    for index, side in enumerate(subscription.iptrunk.iptrunk_sides):
-        nbclient = NetboxClient()
-        if get_router_vendor(side.iptrunk_side_node.owner_subscription_id) == RouterVendor.NOKIA:
-            for interface in side.iptrunk_side_ae_members:
-                if any(
-                    ae_member.get("interface_name") == interface.interface_name
-                    for ae_member in previous_ae_members[str(index)]
-                ):
-                    continue
-                nbclient.allocate_interface(
-                    device_name=side.iptrunk_side_node.router_fqdn,
-                    iface_name=interface.interface_name,
-                )
-            # detach the old interfaces from lag
-            nbclient.detach_interfaces_from_lag(
-                device_name=side.iptrunk_side_node.router_fqdn,
-                lag_name=side.iptrunk_side_ae_iface,
-            )
+    _netbox_allocate_interfaces(subscription.iptrunk.iptrunk_sides[0], previous_ae_members[0])
 
-    return {"subscription": subscription}
+
+@step("Netbox: Allocate side B interfaces")
+def allocate_interfaces_in_netbox_side_b(subscription: Iptrunk, previous_ae_members: list[list[dict]]) -> None:
+    """Allocate the :term:`LAG` interface on side B in Netbox.
+
+    Attach the :term:`LAG` interface to the physical interface detach old one from the :term:`LAG`.
+    """
+    _netbox_allocate_interfaces(subscription.iptrunk.iptrunk_sides[1], previous_ae_members[1])
 
 
 @workflow(
@@ -343,15 +377,29 @@ def modify_trunk_interface() -> StepList:
     * Provision the updated version of the IP trunk, first as a dry run
     * Allocate the reserved interfaces in Netbox
     """
+    side_a_is_nokia = conditional(
+        lambda state: get_router_vendor(
+            state["subscription"]["iptrunk"]["iptrunk_sides"][0]["iptrunk_side_node"]["owner_subscription_id"]
+        )
+        == RouterVendor.NOKIA
+    )
+    side_b_is_nokia = conditional(
+        lambda state: get_router_vendor(
+            state["subscription"]["iptrunk"]["iptrunk_sides"][1]["iptrunk_side_node"]["owner_subscription_id"]
+        )
+        == RouterVendor.NOKIA
+    )
     return (
         init
         >> store_process_subscription(Target.MODIFY)
         >> unsync
         >> modify_iptrunk_subscription
-        >> update_interfaces_in_netbox
+        >> side_a_is_nokia(netbox_update_interfaces_side_a)
+        >> side_b_is_nokia(netbox_update_interfaces_side_b)
         >> pp_interaction(provision_ip_trunk_iface_dry)
         >> pp_interaction(provision_ip_trunk_iface_real)
-        >> allocate_interfaces_in_netbox
+        >> side_a_is_nokia(allocate_interfaces_in_netbox_side_a)
+        >> side_b_is_nokia(allocate_interfaces_in_netbox_side_b)
         >> resync
         >> done
     )