add PE-only steps to update_ibgp_mesh

gso/workflows/router/update_ibgp_mesh.py

@@ -7,7 +7,7 @@ from orchestrator.forms import FormPage
 from orchestrator.forms.validators import Label
 from orchestrator.targets import Target
 from orchestrator.types import FormGenerator, State, SubscriptionLifecycle, UUIDstr
-from orchestrator.workflow import StepList, begin, done, inputstep, step, workflow
+from orchestrator.workflow import StepList, begin, conditional, done, inputstep, 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 ConfigDict, model_validator
@@ -20,6 +20,22 @@ from gso.services.subscriptions import get_trunks_that_terminate_on_router
 from gso.utils.helpers import generate_inventory_for_active_routers
 from gso.utils.types.snmp import SNMPVersion
 from gso.utils.types.tt_number import TTNumber
+from gso.utils.workflow_steps import (
+    add_all_p_to_pe_dry,
+    add_all_p_to_pe_real,
+    add_pe_mesh_to_pe_dry,
+    add_pe_mesh_to_pe_real,
+    add_pe_to_all_p_dry,
+    add_pe_to_all_p_real,
+    add_pe_to_pe_mesh_dry,
+    add_pe_to_pe_mesh_real,
+    check_l3_services,
+    check_pe_ibgp,
+    update_sdp_mesh_dry,
+    update_sdp_mesh_real,
+    update_sdp_single_pe_dry,
+    update_sdp_single_pe_real,
+)
 
 
 def initial_input_form_generator(subscription_id: UUIDstr) -> FormGenerator:
@@ -48,7 +64,9 @@ def initial_input_form_generator(subscription_id: UUIDstr) -> FormGenerator:
 
     user_input = yield AddBGPSessionForm
 
-    return user_input.model_dump()
+    return user_input.model_dump() | {
+        "router_role": subscription.router.router_role,
+    }
 
 
 @step("[DRY RUN] Add P router to iBGP mesh")
@@ -201,15 +219,32 @@ def update_ibgp_mesh() -> StepList:
     * Add the new P-router to LibreNMS.
     * Update the subscription model.
     """
+    router_is_pe = conditional(lambda state: state["router_role"] == RouterRole.PE)
+    router_is_p = conditional(lambda state: state["router_role"] == RouterRole.P)
+
     return (
         begin
         >> store_process_subscription(Target.MODIFY)
         >> unsync
-        >> lso_interaction(add_p_to_mesh_dry)
-        >> lso_interaction(add_p_to_mesh_real)
-        >> lso_interaction(add_all_pe_to_p_dry)
-        >> lso_interaction(add_all_pe_to_p_real)
-        >> lso_interaction(check_ibgp_session)
+        >> router_is_p(lso_interaction(add_p_to_mesh_dry))
+        >> router_is_p(lso_interaction(add_p_to_mesh_real))
+        >> router_is_p(lso_interaction(add_all_pe_to_p_dry))
+        >> router_is_p(lso_interaction(add_all_pe_to_p_real))
+        >> router_is_p(lso_interaction(check_ibgp_session))
+        >> router_is_pe(lso_interaction(add_pe_mesh_to_pe_dry))
+        >> router_is_pe(lso_interaction(add_pe_mesh_to_pe_real))
+        >> router_is_pe(lso_interaction(add_pe_to_pe_mesh_dry))
+        >> router_is_pe(lso_interaction(add_pe_to_pe_mesh_real))
+        >> router_is_pe(lso_interaction(add_all_p_to_pe_dry))
+        >> router_is_pe(lso_interaction(add_all_p_to_pe_real))
+        >> router_is_pe(lso_interaction(add_pe_to_all_p_dry))
+        >> router_is_pe(lso_interaction(add_pe_to_all_p_real))
+        >> router_is_pe(lso_interaction(update_sdp_single_pe_dry))
+        >> router_is_pe(lso_interaction(update_sdp_single_pe_real))
+        >> router_is_pe(lso_interaction(update_sdp_mesh_dry))
+        >> router_is_pe(lso_interaction(update_sdp_mesh_real))
+        >> router_is_pe(lso_interaction(check_pe_ibgp))
+        >> router_is_pe(lso_interaction(check_l3_services))
         >> add_device_to_librenms
         >> prompt_insert_in_radius
         >> prompt_radius_login