Add support for modifying the IAS flavour in creation and modification workflows

gso/workflows/l3_core_service/base_modify_l3_core_service.py

@@ -214,7 +214,7 @@ def initial_input_form_generator(subscription_id: UUIDstr) -> FormGenerator:
             user_input = yield AddAccessPortForm
             return {
                 "operation": initial_input.operation,
-                "added_access_port": user_input,
+                "added_access_port": user_input.model_dump(),
             }
 
         case Operation.REMOVE:
@@ -357,7 +357,7 @@ def initial_input_form_generator(subscription_id: UUIDstr) -> FormGenerator:
             return {
                 "operation": initial_input.operation,
                 "modified_access_port": user_input.access_port,
-                "modified_sbp": binding_port_input_form,
+                "modified_sbp": binding_port_input_form.model_dump(),
             }
 
         case _:

gso/workflows/l3_core_service/ias/create_ias.py

@@ -26,6 +26,7 @@ from gso.workflows.l3_core_service.base_create_l3_core_service import (
 from gso.workflows.l3_core_service.base_create_l3_core_service import (
     initial_input_form_generator as base_initial_input_form_generator,
 )
+from gso.workflows.l3_core_service.ias.shared import update_ias_subscription_model
 
 
 def initial_input_form_generator(product_name: str) -> FormGenerator:
@@ -61,6 +62,7 @@ def create_ias() -> StepList:
         >> create_subscription
         >> store_process_subscription(Target.CREATE)
         >> initialize_subscription
+        >> update_ias_subscription_model
         >> start_moodi()
         >> lso_interaction(provision_sbp_dry)
         >> lso_interaction(provision_sbp_real)

gso/workflows/l3_core_service/ias/create_imported_ias.py

@@ -19,6 +19,7 @@ from gso.workflows.l3_core_service.base_create_imported_l3_core_service import (
 from gso.workflows.l3_core_service.base_create_imported_l3_core_service import (
     initialize_subscription,
 )
+from gso.workflows.l3_core_service.ias.shared import update_ias_subscription_model
 
 
 def initial_input_form_generator() -> FormGenerator:
@@ -53,6 +54,7 @@ def create_imported_ias() -> StepList:
     return (
         begin
         >> create_subscription
+        >> update_ias_subscription_model
         >> store_process_subscription(Target.CREATE)
         >> initialize_subscription
         >> set_status(SubscriptionLifecycle.ACTIVE)

gso/workflows/l3_core_service/ias/modify_ias.py

 """Modification workflow for a IAS subscription."""
 
 from orchestrator import begin, conditional, done, workflow
+from orchestrator.domain import SubscriptionModel
 from orchestrator.targets import Target
 from orchestrator.workflow import StepList
 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 gso.products.product_blocks.ias import IASFlavor
 from gso.workflows.l3_core_service.base_modify_l3_core_service import (
     Operation,
     create_new_sbp,
@@ -13,11 +17,28 @@ from gso.workflows.l3_core_service.base_modify_l3_core_service import (
     modify_existing_sbp,
     remove_old_sbp,
 )
+from gso.workflows.l3_core_service.ias.shared import update_ias_subscription_model
+
+
+def modify_ias_input_form_generator(subscription_id: UUIDstr) -> FormGenerator:
+    """Initial form generator for modifying the custom attributes of an existing IAS subscription."""
+    initial_generator = initial_input_form_generator(subscription_id)
+    initial_user_input = yield from initial_generator
+
+    subscription = SubscriptionModel.from_subscription(subscription_id)
+
+    # Additional IAS step
+    class IASExtraForm(FormPage):
+        # TODO: remove type hint workaround
+        ias_flavor: IASFlavor | str = subscription.ias.ias_flavor  # type: ignore[attr-defined]
+
+    ias_extra = yield IASExtraForm
+    return initial_user_input | ias_extra.model_dump()
 
 
 @workflow(
     "Modify IAS",
-    initial_input_form=wrap_modify_initial_input_form(initial_input_form_generator),
+    initial_input_form=wrap_modify_initial_input_form(modify_ias_input_form_generator),
     target=Target.MODIFY,
 )
 def modify_ias() -> StepList:
@@ -30,6 +51,7 @@ def modify_ias() -> StepList:
         begin
         >> store_process_subscription(Target.MODIFY)
         >> unsync
+        >> update_ias_subscription_model
         >> access_port_is_added(create_new_sbp)
         >> access_port_is_removed(remove_old_sbp)
         >> access_port_is_modified(modify_existing_sbp)

gso/workflows/l3_core_service/ias/shared.py

+"""Shared logic for IAS service workflows."""
+
+from orchestrator import step
+from orchestrator.domain import SubscriptionModel
+from pydantic_forms.types import State
+
+from gso.products.product_blocks.ias import IASFlavor
+
+
+@step("Update IAS-specific attributes")
+def update_ias_subscription_model(subscription: SubscriptionModel, ias_flavor: IASFlavor) -> State:
+    """Update the subscription model of an IAS subscription with a new IAS flavour."""
+    subscription.ias.ias_flavor = ias_flavor  # type: ignore[attr-defined]
+
+    return {"subscription": subscription}

test/workflows/l3_core_service/test_create_imported_l3_core_service.py

@@ -14,7 +14,7 @@ from test.workflows import assert_complete, extract_state, run_workflow
 @pytest.mark.parametrize("product_name", L3_PRODUCT_NAMES)
 def test_create_imported_l3_core_service_success(faker, partner_factory, edge_port_subscription_factory, product_name):
     extra_ias_data = {
-        "ias_flavor": IASFlavor.IAS_PS_OPT_OUT,
+        "ias_flavor": IASFlavor.IASGWS,
     }
     creation_form_input_data = {
         "partner": partner_factory()["name"],
@@ -83,3 +83,6 @@ def test_create_imported_l3_core_service_success(faker, partner_factory, edge_po
     assert_complete(result)
     assert subscription.status == SubscriptionLifecycle.ACTIVE
     assert subscription.product.name == PRODUCT_IMPORTED_MAP[product_name]
+
+    if product_name == ProductName.IAS:
+        assert subscription.ias.ias_flavor == IASFlavor.IASGWS

test/workflows/l3_core_service/test_create_l3_core_service.py

@@ -106,6 +106,7 @@ def test_create_l3_core_service_success(
     assert_complete(result)
     state = extract_state(result)
     subscription = SubscriptionModel.from_subscription(state["subscription_id"])
+
     assert subscription.product.name == product_name
     assert mock_lso_client.call_count == lso_interaction_count + 1
     assert subscription.status == SubscriptionLifecycle.ACTIVE
@@ -116,3 +117,6 @@ def test_create_l3_core_service_success(
     )
     assert subscription.l3_core.ap_list[0].sbp.gs_id == "GS-12345"
     assert mock_sharepoint_client.call_count == 1
+
+    if product_name == ProductName.IAS:
+        assert subscription.ias.ias_flavor == IASFlavor.IASGWS

test/workflows/l3_core_service/test_modify_l3_core_service.py

 import pytest
 from orchestrator.domain import SubscriptionModel
 
+from gso.products import ProductName
 from gso.products.product_blocks.bgp_session import IPFamily
+from gso.products.product_blocks.ias import IASFlavor
 from gso.utils.shared_enums import APType
 from gso.workflows.l3_core_service.base_modify_l3_core_service import Operation
 from gso.workflows.l3_core_service.shared import L3_MODIFICATION_WF_MAP, L3_PRODUCT_NAMES
@@ -19,6 +21,12 @@ def test_modify_l3_core_service_remove_edge_port_success(faker, l3_core_service_
         {"access_port": str(access_port.subscription_instance_id)},
     ]
 
+    if product_name == ProductName.IAS:
+        extra_ias_data = {
+            "ias_flavor": IASFlavor.IASGWS,
+        }
+        input_form_data.append(extra_ias_data)
+
     result, _, _ = run_workflow(L3_MODIFICATION_WF_MAP[product_name], input_form_data)
 
     state = extract_state(result)
@@ -26,6 +34,8 @@ def test_modify_l3_core_service_remove_edge_port_success(faker, l3_core_service_
     ap_list = subscription.l3_core.ap_list
     assert len(ap_list) == 1
     assert ap_list[0].ap_type == APType.BACKUP
+    if product_name == ProductName.IAS:
+        assert subscription.ias.ias_flavor == IASFlavor.IASGWS
 
 
 @pytest.mark.parametrize("product_name", L3_PRODUCT_NAMES)
@@ -70,6 +80,12 @@ def test_modify_l3_core_service_add_new_edge_port_success(
         },
     ]
 
+    if product_name == ProductName.IAS:
+        extra_ias_data = {
+            "ias_flavor": IASFlavor.IASGWS,
+        }
+        input_form_data.append(extra_ias_data)
+
     result, _, _ = run_workflow(L3_MODIFICATION_WF_MAP[product_name], input_form_data)
 
     state = extract_state(result)
@@ -84,6 +100,8 @@ def test_modify_l3_core_service_add_new_edge_port_success(
     assert str(new_ap.sbp.ipv6_address) == input_form_data[2]["ipv6_address"]
     assert new_ap.sbp.ipv6_mask == input_form_data[2]["ipv6_mask"]
     assert len(ap_list) == 3
+    if product_name == ProductName.IAS:
+        assert subscription.ias.ias_flavor == IASFlavor.IASGWS
 
 
 @pytest.fixture()
@@ -145,6 +163,12 @@ def test_modify_l3_core_service_modify_edge_port_success(
         {**new_sbp_data},
     ]
 
+    if product_name == ProductName.IAS:
+        extra_ias_data = {
+            "ias_flavor": IASFlavor.IASGWS,
+        }
+        input_form_data.append(extra_ias_data)
+
     result, _, _ = run_workflow(L3_MODIFICATION_WF_MAP[product_name], input_form_data)
 
     state = extract_state(result)
@@ -193,3 +217,6 @@ def test_modify_l3_core_service_modify_edge_port_success(
     assert access_port.sbp.v6_bfd_settings.bfd_interval_rx == new_sbp_data["v6_bfd_interval_rx"]
     assert access_port.sbp.v6_bfd_settings.bfd_interval_tx == new_sbp_data["v6_bfd_interval_tx"]
     assert access_port.sbp.v6_bfd_settings.bfd_multiplier == new_sbp_data["v6_bfd_multiplier"]
+
+    if product_name == ProductName.IAS:
+        assert subscription.ias.ias_flavor == IASFlavor.IASGWS