Update iBGP mesh workflow and unit test

gso/workflows/router/update_ibgp_mesh.py

@@ -210,8 +210,8 @@ 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)
+    router_is_pe = conditional(lambda state: state["subscription"]["router"]["router_role"] == RouterRole.PE)
+    router_is_p = conditional(lambda state: state["subscription"]["router"]["router_role"] == RouterRole.P)
 
     return (
         begin

test/workflows/router/test_update_ibgp_mesh.py

@@ -18,6 +18,7 @@ from test.workflows import (
 
 
 @pytest.mark.parametrize("trunk_status", [SubscriptionLifecycle.PROVISIONING, SubscriptionLifecycle.ACTIVE])
+@pytest.mark.parametrize("router_role", [RouterRole.P, RouterRole.PE])
 @pytest.mark.workflow()
 @patch("gso.services.lso_client._send_request")
 @patch("gso.workflows.router.update_ibgp_mesh.librenms_client.LibreNMSClient.add_device")
@@ -26,13 +27,27 @@ def test_update_ibgp_mesh_success(
     mock_librenms_device_exists,
     mock_librenms_add_device,
     mock_execute_playbook,
+    router_role,
     trunk_status,
     iptrunk_subscription_factory,
+    iptrunk_side_subscription_factory,
+    nokia_router_subscription_factory,
     data_config_filename,
     faker,
 ):
     mock_librenms_device_exists.return_value = False
-    ip_trunk = Iptrunk.from_subscription(iptrunk_subscription_factory(status=trunk_status))
+    side_a = iptrunk_side_subscription_factory(
+        iptrunk_side_node=nokia_router_subscription_factory(router_role=router_role)
+    )
+    side_b = iptrunk_side_subscription_factory()
+    #  Add some extra devices to the network.
+    nokia_router_subscription_factory(router_role=RouterRole.P)
+    nokia_router_subscription_factory()
+
+    ip_trunk = Iptrunk.from_subscription(
+        iptrunk_subscription_factory(status=trunk_status, iptrunk_sides=[side_a, side_b])
+    )
+    callback_step_count = 5 if router_role == RouterRole.P else 14
     ibgp_mesh_input_form_data = {
         "subscription_id": ip_trunk.iptrunk.iptrunk_sides[0].iptrunk_side_node.owner_subscription_id
     }
@@ -40,7 +55,7 @@ def test_update_ibgp_mesh_success(
         "update_ibgp_mesh", [ibgp_mesh_input_form_data, {"tt_number": faker.tt_number()}]
     )
 
-    for _ in range(5):
+    for _ in range(callback_step_count):
         result, step_log = assert_lso_interaction_success(result, process_stat, step_log)
 
     # Handle two consecutive user input steps
@@ -50,7 +65,7 @@ def test_update_ibgp_mesh_success(
 
     state = extract_state(result)
 
-    assert mock_execute_playbook.call_count == 5
+    assert mock_execute_playbook.call_count == callback_step_count
     assert mock_librenms_add_device.call_count == 1
     assert result.status == StepStatus.COMPLETE
     assert state["subscription"]["router"]["router_access_via_ts"] is False