Added modify connection strategy workflow

gso/migrations/versions/2023-11-21_e8378fbcfbf3_add_initial_products.py

@@ -11,7 +11,7 @@ from alembic import op
 # revision identifiers, used by Alembic.
 revision = 'e8378fbcfbf3'
 down_revision = 'da5c9f4cce1c'
-branch_labels = ('data',)
+branch_labels = None
 depends_on = None
 
 

gso/migrations/versions/2024-03-21_a2cd3f2e6d7a_modify_connection_streategy_workflow.py

+"""Modify connection streategy workflow..
+
+Revision ID: a2cd3f2e6d7a
+Revises:
+Create Date: 2024-03-21 16:05:59.043106
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = 'a2cd3f2e6d7a'
+down_revision = None
+branch_labels = None
+depends_on = 'd61c0f92da1e'
+
+
+from orchestrator.migrations.helpers import create_workflow, delete_workflow
+
+new_workflows = [
+    {
+        "name": "modify_connection_strategy",
+        "target": "MODIFY",
+        "description": "Modify connection strategy",
+        "product_type": "Router"
+    }
+]
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    for workflow in new_workflows:
+        create_workflow(conn, workflow)
+
+
+def downgrade() -> None:
+    conn = op.get_bind()
+    for workflow in new_workflows:
+        delete_workflow(conn, workflow["name"])

gso/utils/shared_enums.py

@@ -19,3 +19,10 @@ class PortNumber(ConstrainedInt):
 
     gt = 0
     le = 49151
+
+
+class ConnectionStrategy(strEnum):
+    """An enumerator for the connection Strategies."""
+
+    IN_BAND = "IN BAND"
+    OUT_OF_BAND = "OUT OF BAND"

gso/workflows/__init__.py

@@ -26,6 +26,7 @@ LazyWorkflowInstance("gso.workflows.router.create_router", "create_router")
 LazyWorkflowInstance("gso.workflows.router.redeploy_base_config", "redeploy_base_config")
 LazyWorkflowInstance("gso.workflows.router.terminate_router", "terminate_router")
 LazyWorkflowInstance("gso.workflows.router.update_ibgp_mesh", "update_ibgp_mesh")
+LazyWorkflowInstance("gso.workflows.router.modify_connection_strategy", "modify_connection_strategy")
 LazyWorkflowInstance("gso.workflows.site.create_site", "create_site")
 LazyWorkflowInstance("gso.workflows.site.modify_site", "modify_site")
 LazyWorkflowInstance("gso.workflows.site.terminate_site", "terminate_site")

gso/workflows/router/modify_connection_strategy.py

+"""Modify connection strategy workflow. Flipping the connection between in-band to out-of-band."""
+from orchestrator.forms import FormPage
+from orchestrator.targets import Target
+from orchestrator.types import FormGenerator, State, UUIDstr
+from orchestrator.workflow import StepList, 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 gso.products.product_types.router import Router
+from gso.utils.shared_enums import ConnectionStrategy
+
+
+def initial_input_form_generator(subscription_id: UUIDstr) -> FormGenerator:
+    """Initial form with a choice field to modify the connection strategy."""
+    subscription = Router.from_subscription(subscription_id)
+
+    current_connection_strategy = (
+        ConnectionStrategy.OUT_OF_BAND
+        if subscription.router.router_access_via_ts
+        else ConnectionStrategy.IN_BAND
+    )
+
+    class ModifyConnectionStrategyForm(FormPage):
+        class Config:
+            title = f"Modify the connection strategy of {subscription.router.router_fqdn}."
+
+        connection_strategy: ConnectionStrategy = current_connection_strategy
+
+    user_input = yield ModifyConnectionStrategyForm
+
+    return user_input.dict()
+
+
+@step("Update subscription model")
+def update_subscription_model(subscription: Router, connection_strategy: str) -> State:
+    """Update the database model to update connection strategy."""
+    subscription.router.router_access_via_ts = False if connection_strategy == ConnectionStrategy.IN_BAND else True
+
+    return {"subscription": subscription}
+
+
+@workflow(
+    "Modify connection strategy",
+    initial_input_form=wrap_modify_initial_input_form(initial_input_form_generator),
+    target=Target.MODIFY,
+)
+def modify_connection_strategy() -> StepList:
+    """Modify the connection strategy."""
+    return (
+            init
+            >> store_process_subscription(Target.MODIFY)
+            >> unsync
+            >> update_subscription_model
+            >> resync
+            >> done
+    )

test/workflows/router/test_modify_connection_stratey.py

+import pytest
+
+from gso.products import Router
+from gso.utils.shared_enums import ConnectionStrategy
+from test.workflows import assert_complete, run_workflow
+
+
+@pytest.mark.workflow()
+def test_modify_connection_strategy(responses, nokia_router_subscription_factory):
+    subscription_id = nokia_router_subscription_factory(router_access_via_ts=True)
+    subscription = Router.from_subscription(subscription_id)
+    assert subscription.router.router_access_via_ts is True
+    form_data = [
+        {"subscription_id": subscription_id},
+        {
+            "connection_strategy": ConnectionStrategy.IN_BAND
+        },
+    ]
+    result, _, _ = run_workflow("modify_connection_strategy", form_data)
+    assert_complete(result)
+
+    # Fetch the updated subscription after running the workflow
+    updated_subscription = Router.from_subscription(subscription_id)
+    assert updated_subscription.status == "active"
+    assert updated_subscription.router.router_access_via_ts is False
+