Split up validation workflow for juniper and nokia routers

gso/schedules/__init__.py

 """Tasks that are scheduled to run periodically in :term:`GSO`."""
+
+TT_NUMBER_ZERO = "TT#0000000000000000"

gso/workflows/router/validate_router.py

@@ -2,29 +2,29 @@
 
 from orchestrator.targets import Target
 from orchestrator.utils.errors import ProcessFailureError
-from orchestrator.workflow import StepList, begin, done, step, workflow
+from orchestrator.workflow import StepList, begin, conditional, done, step, workflow
 from orchestrator.workflows.steps import resync, store_process_subscription, unsync
 from orchestrator.workflows.utils import wrap_modify_initial_input_form
-from pydantic_forms.core import FormPage
-from pydantic_forms.types import FormGenerator, UUIDstr
+from pydantic_forms.types import State, UUIDstr
 
 from gso.products.product_types.router import Router
+from gso.schedules import TT_NUMBER_ZERO
 from gso.services import infoblox
 from gso.services.librenms_client import LibreNMSClient
 from gso.services.lso_client import anonymous_lso_interaction
 from gso.services.netbox_client import NetboxClient
+from gso.utils.shared_enums import Vendor
 from gso.utils.workflow_steps import deploy_base_config_dry, detect_configuration_drift
 from gso.workflows.router.update_ibgp_mesh import add_p_to_mesh_dry
 
-TT_NUMBER_ZERO = "TT#0000000000000000"
 
+@step("Prepare required keys in state")
+def prepare_state(subscription_id: UUIDstr) -> State:
+    """Add required keys to the state for the workflow to run successfully.
 
-def _seed_initial_state() -> FormGenerator:
-    class InputForm(FormPage):
-        subscription_id: UUIDstr
-
-    user_input = yield InputForm
-    router = Router.from_subscription(user_input["subscription_id"])
+    This includes the router product itself and a fake TT number for running playbooks.
+    """
+    router = Router.from_subscription(subscription_id)
 
     return {"subscription": router, "tt_number": TT_NUMBER_ZERO}
 
@@ -64,9 +64,7 @@ def verify_librenms_entry(subscription: Router) -> None:
 
 
 @workflow(
-    "Validate router configuration",
-    target=Target.SYSTEM,
-    initial_input_form=wrap_modify_initial_input_form(_seed_initial_state),
+    "Validate router configuration", target=Target.SYSTEM, initial_input_form=wrap_modify_initial_input_form(None)
 )
 def validate_router() -> StepList:
     """Validate an existing, active Router subscription.
@@ -77,9 +75,13 @@ def validate_router() -> StepList:
     * Redeploy base config to verify the configuration is intact.
     * Validate configuration of the iBGP mesh
     """
+    is_juniper_router = conditional(lambda state: state["subscription"]["router"]["vendor"] == Vendor.JUNIPER)
+
     return (
         begin
         >> store_process_subscription(Target.SYSTEM)
+        >> prepare_state
+        >> is_juniper_router(done)
         >> unsync
         >> verify_ipam_loopback
         >> verify_netbox_entry

test/workflows/__init__.py

@@ -282,6 +282,12 @@ def user_accept_and_assert_suspended(process_stat, step_log, extra_data=None):
     return result, step_log
 
 
+def assert_lso_success(result: Process, process_stat: ProcessStat, step_log: list):
+    """Assert a successful LSO execution in a workflow."""
+    assert_awaiting_callback(result)
+    return resume_workflow(process_stat, step_log, input_data=LSO_RESULT_SUCCESS)
+
+
 def assert_lso_interaction_success(result: Process, process_stat: ProcessStat, step_log: list):
     """Assert a successful LSO interaction in a workflow.
 
@@ -290,7 +296,7 @@ def assert_lso_interaction_success(result: Process, process_stat: ProcessStat, s
     waiting for the user to confirm the results received.
     """
     assert_awaiting_callback(result)
-    result, step_log = resume_workflow(process_stat, step_log, input_data=LSO_RESULT_SUCCESS)
+    result, step_log = assert_lso_success(result, process_stat, step_log)
     assert_suspended(result)
 
     return resume_workflow(process_stat, step_log, input_data=USER_CONFIRM_EMPTY_FORM)

test/workflows/router/test_validate_router.py

@@ -6,27 +6,29 @@ from infoblox_client import objects
 from gso.products.product_types.router import Router
 from test.workflows import (
     assert_complete,
-    assert_lso_interaction_success,
+    assert_lso_success,
     extract_state,
     run_workflow,
 )
 
 
 @pytest.mark.workflow()
-@pytest.mark.parametrize("product_id", ["nokia_router_subscription_factory", "juniper_router_subscription_factory"])
 @patch("gso.services.infoblox.find_host_by_fqdn")
 @patch("gso.services.lso_client.execute_playbook")
-def test_validate_router_success(
-    mock_validate_router,
+@patch("gso.services.netbox_client.NetboxClient.get_device_by_name")
+@patch("gso.services.librenms_client.LibreNMSClient.get_device")
+def test_validate_nokia_router_success(
+    mock_get_librenms_device,
+    mock_get_device_by_name,
+    mock_execute_playbook,
     mock_find_host_by_fqdn,
-    product_id,
+    nokia_router_subscription_factory,
     faker,
     data_config_filename,
-    request,
     geant_partner,
 ):
     #  Run workflow
-    subscription_id = request.getfixturevalue(product_id)(partner=geant_partner)
+    subscription_id = nokia_router_subscription_factory()
     mock_fqdn = Router.from_subscription(subscription_id).router.router_fqdn
     mock_v4 = faker.ipv4()
     mock_find_host_by_fqdn.return_value = objects.HostRecord(
@@ -51,7 +53,8 @@ def test_validate_router_success(
     state = extract_state(result)
     subscription_id = state["subscription_id"]
 
-    result, step_log = assert_lso_interaction_success(result, process_stat, step_log)
+    for _ in range(2):
+        result, step_log = assert_lso_success(result, process_stat, step_log)
 
     assert_complete(result)
 
@@ -59,6 +62,29 @@ def test_validate_router_success(
     subscription = Router.from_subscription(subscription_id)
 
     assert subscription.status == "active"
-    assert mock_validate_router.call_count == 1
+    assert mock_execute_playbook.call_count == 2
     assert mock_find_host_by_fqdn.call_count == 1
+    assert mock_get_device_by_name.call_count == 1
+    assert mock_get_librenms_device.call_count == 1
+    assert "ipam_warning" not in state
+
+
+@pytest.mark.workflow()
+def test_validate_juniper_router_success(
+    juniper_router_subscription_factory,
+    faker,
+    data_config_filename,
+    geant_partner,
+):
+    #  Run workflow
+    subscription_id = juniper_router_subscription_factory()
+
+    initial_router_data = [{"subscription_id": subscription_id}]
+    result, _, _ = run_workflow("validate_router", initial_router_data)
+
+    assert_complete(result)
+    state = extract_state(result)
+    subscription = Router.from_subscription(subscription_id)
+
+    assert subscription.status == "active"
     assert "ipam_warning" not in state