initial mgmt group stage data classes extend the management group data class

atst/domain/csp/cloud/azure_cloud_provider.py

@@ -25,11 +25,12 @@ from .models import (
     BillingProfileVerificationCSPPayload,
     BillingProfileVerificationCSPResult,
     CostManagementQueryCSPResult,
+    InitialMgmtGroupCSPPayload,
+    InitialMgmtGroupCSPResponse,
+    InitialMgmtGroupVerificationCSPPayload,
+    InitialMgmtGroupVerificationCSPResponse,
     KeyVaultCredentials,
-    ManagementGroupCSPPayload,
     ManagementGroupCSPResponse,
-    ManagementGroupGetCSPPayload,
-    ManagementGroupGetCSPResponse,
     ProductPurchaseCSPPayload,
     ProductPurchaseCSPResult,
     ProductPurchaseVerificationCSPPayload,
@@ -216,7 +217,7 @@ class AzureCloudProvider(CloudProviderInterface):
 
         return ApplicationCSPResult(**response)
 
-    def create_initial_mgmt_group(self, payload: ManagementGroupCSPPayload):
+    def create_initial_mgmt_group(self, payload: InitialMgmtGroupCSPPayload):
         creds = self._source_creds(payload.tenant_id)
         credentials = self._get_credential_obj(
             {
@@ -230,10 +231,10 @@ class AzureCloudProvider(CloudProviderInterface):
             credentials, payload.management_group_name, payload.display_name,
         )
 
-        return ManagementGroupCSPResponse(**response)
+        return InitialMgmtGroupCSPResponse(**response)
 
     def create_initial_mgmt_group_verification(
-        self, payload: ManagementGroupGetCSPPayload
+        self, payload: InitialMgmtGroupVerificationCSPPayload
     ):
         creds = self._source_creds(payload.tenant_id)
         credentials = self._get_credential_obj(
@@ -246,7 +247,7 @@ class AzureCloudProvider(CloudProviderInterface):
         )
 
         response = self._get_management_group(credentials, payload.tenant_id,)
-        return ManagementGroupGetCSPResponse(**response.result())
+        return InitialMgmtGroupVerificationCSPResponse(**response.result())
 
     def _create_management_group(
         self, credentials, management_group_id, display_name, parent_id=None,

atst/domain/csp/cloud/mock_cloud_provider.py

@@ -25,10 +25,10 @@ from .models import (
     BillingProfileTenantAccessCSPResult,
     BillingProfileVerificationCSPPayload,
     BillingProfileVerificationCSPResult,
-    ManagementGroupCSPPayload,
+    InitialMgmtGroupCSPPayload,
-    ManagementGroupCSPResponse,
+    InitialMgmtGroupCSPResponse,
-    ManagementGroupGetCSPPayload,
+    InitialMgmtGroupVerificationCSPPayload,
-    ManagementGroupGetCSPResponse,
+    InitialMgmtGroupVerificationCSPResponse,
     CostManagementQueryCSPResult,
     CostManagementQueryProperties,
     ProductPurchaseCSPPayload,
@@ -329,23 +329,23 @@ class MockCloudProvider(CloudProviderInterface):
             }
         )
 
-    def create_initial_mgmt_group(self, payload: ManagementGroupCSPPayload):
+    def create_initial_mgmt_group(self, payload: InitialMgmtGroupCSPPayload):
         self._maybe_raise(self.NETWORK_FAILURE_PCT, self.NETWORK_EXCEPTION)
         self._maybe_raise(self.SERVER_FAILURE_PCT, self.SERVER_EXCEPTION)
         self._maybe_raise(self.UNAUTHORIZED_RATE, self.AUTHORIZATION_EXCEPTION)
 
-        return ManagementGroupCSPResponse(
+        return InitialMgmtGroupCSPResponse(
-            id=f"{AZURE_MGMNT_PATH}{payload.management_group_name}"
+            id=f"{AZURE_MGMNT_PATH}{payload.management_group_name}",
         )
 
     def create_initial_mgmt_group_verification(
-        self, payload: ManagementGroupGetCSPPayload
+        self, payload: InitialMgmtGroupVerificationCSPPayload
     ):
         self._maybe_raise(self.NETWORK_FAILURE_PCT, self.NETWORK_EXCEPTION)
         self._maybe_raise(self.SERVER_FAILURE_PCT, self.SERVER_EXCEPTION)
         self._maybe_raise(self.UNAUTHORIZED_RATE, self.AUTHORIZATION_EXCEPTION)
 
-        return ManagementGroupGetCSPResponse(
+        return InitialMgmtGroupVerificationCSPResponse(
             **dict(
                 id="Test Id"
                 # id=f"{AZURE_MGMNT_PATH}{payload.management_group_name}"

atst/domain/csp/cloud/models.py

@@ -367,6 +367,22 @@ class ApplicationCSPResult(ManagementGroupCSPResponse):
     pass
 
 
+class InitialMgmtGroupCSPPayload(ManagementGroupCSPPayload):
+    pass
+
+
+class InitialMgmtGroupCSPResponse(ManagementGroupCSPResponse):
+    pass
+
+
+class InitialMgmtGroupVerificationCSPPayload(ManagementGroupGetCSPPayload):
+    pass
+
+
+class InitialMgmtGroupVerificationCSPResponse(ManagementGroupGetCSPResponse):
+    pass
+
+
 class KeyVaultCredentials(BaseModel):
     root_sp_client_id: Optional[str]
     root_sp_key: Optional[str]

atst/models/portfolio_state_machine.py

@@ -155,37 +155,37 @@ class PortfolioStateMachine(
             print(exc.json())
             app.logger.info(payload)
             self.fail_stage(stage)
+        else:
+            # TODO: Determine best place to do this, maybe @reconstructor
+            self.csp = app.csp.cloud
 
-        # TODO: Determine best place to do this, maybe @reconstructor
+            try:
-        self.csp = app.csp.cloud
+                func_name = f"create_{stage}"
+                response = getattr(self.csp, func_name)(payload_data)
+                if self.portfolio.csp_data is None:
+                    self.portfolio.csp_data = {}
+                self.portfolio.csp_data.update(response.dict())
+                db.session.add(self.portfolio)
+                db.session.commit()
+            except PydanticValidationError as exc:
+                app.logger.error(
+                    f"Failed to cast response to valid result class {self.__repr__()}:",
+                    exc_info=1,
+                )
+                app.logger.info(exc.json())
+                print(exc.json())
+                app.logger.info(payload_data)
+                # TODO: Ensure that failing the stage does not preclude a Celery retry
+                self.fail_stage(stage)
+            # TODO: catch and handle general CSP exception here
+            except (ConnectionException, UnknownServerException) as exc:
+                app.logger.error(
+                    f"CSP api call. Caught exception for {self.__repr__()}.", exc_info=1,
+                )
+                # TODO: Ensure that failing the stage does not preclude a Celery retry
+                self.fail_stage(stage)
 
-        try:
+            self.finish_stage(stage)
-            func_name = f"create_{stage}"
-            response = getattr(self.csp, func_name)(payload_data)
-            if self.portfolio.csp_data is None:
-                self.portfolio.csp_data = {}
-            self.portfolio.csp_data.update(response.dict())
-            db.session.add(self.portfolio)
-            db.session.commit()
-        except PydanticValidationError as exc:
-            app.logger.error(
-                f"Failed to cast response to valid result class {self.__repr__()}:",
-                exc_info=1,
-            )
-            app.logger.info(exc.json())
-            print(exc.json())
-            app.logger.info(payload_data)
-            # TODO: Ensure that failing the stage does not preclude a Celery retry
-            self.fail_stage(stage)
-        # TODO: catch and handle general CSP exception here
-        except (ConnectionException, UnknownServerException) as exc:
-            app.logger.error(
-                f"CSP api call. Caught exception for {self.__repr__()}.", exc_info=1,
-            )
-            # TODO: Ensure that failing the stage does not preclude a Celery retry
-            self.fail_stage(stage)
-
-        self.finish_stage(stage)
 
     def is_csp_data_valid(self, event):
         """

tests/domain/cloud/test_azure_csp.py

@@ -22,10 +22,10 @@ from atst.domain.csp.cloud.models import (
     BillingProfileTenantAccessCSPResult,
     BillingProfileVerificationCSPPayload,
     BillingProfileVerificationCSPResult,
-    ManagementGroupCSPPayload,
+    InitialMgmtGroupCSPPayload,
-    ManagementGroupCSPResponse,
+    InitialMgmtGroupCSPResponse,
-    ManagementGroupGetCSPPayload,
+    InitialMgmtGroupVerificationCSPPayload,
-    ManagementGroupGetCSPResponse,
+    InitialMgmtGroupVerificationCSPResponse,
     CostManagementQueryCSPResult,
     ProductPurchaseCSPPayload,
     ProductPurchaseCSPResult,
@@ -116,13 +116,12 @@ def test_create_initial_mgmt_group_succeeds(mock_azure: AzureCloudProvider):
     mock_management_group_create(mock_azure, {"id": "Test Id"})
     mock_azure = mock_get_secret(mock_azure)
 
-    payload = ManagementGroupCSPPayload(
+    payload = InitialMgmtGroupCSPPayload(
         tenant_id="1234",
         display_name=application.name,
         management_group_name=str(uuid4()),
     )
-
+    result: InitialMgmtGroupCSPResponse = mock_azure.create_initial_mgmt_group(payload)
-    result: ManagementGroupCSPResponse = mock_azure.create_initial_mgmt_group(payload)
 
     assert result.id == "Test Id"
 
@@ -136,11 +135,10 @@ def test_create_initial_mgmt_group_verification_succeeds(
 
     management_group_name = str(uuid4())
 
-    payload = ManagementGroupGetCSPPayload(
+    payload = InitialMgmtGroupVerificationCSPPayload(
         tenant_id="1234", management_group_name=management_group_name
     )
-
+    result: InitialMgmtGroupVerificationCSPResponse = mock_azure.create_initial_mgmt_group_verification(
-    result: ManagementGroupGetCSPResponse = mock_azure.create_initial_mgmt_group_verification(
         payload
     )
 

tests/domain/test_portfolio_state_machine.py

@@ -112,7 +112,7 @@ def test_fsm_transition_start(mock_cloud_provider, portfolio: Portfolio):
         FSMStates.ADMIN_ROLE_DEFINITION_CREATED,
         FSMStates.PRINCIPAL_ADMIN_ROLE_CREATED,
         FSMStates.INITIAL_MGMT_GROUP_CREATED,
-        FSMStates.PRODUCT_PURCHASE_VERIFICATION_CREATED,
+        FSMStates.INITIAL_MGMT_GROUP_VERIFICATION_CREATED,
         FSMStates.TENANT_ADMIN_OWNERSHIP_CREATED,
         FSMStates.TENANT_PRINCIPAL_OWNERSHIP_CREATED,
     ]
@@ -133,6 +133,8 @@ def test_fsm_transition_start(mock_cloud_provider, portfolio: Portfolio):
         "user_id": user_id,
         "password": "jklfsdNCVD83nklds2#202",  # pragma: allowlist secret
         "domain_name": domain_name,
+        "display_name" : "mgmt group display name",
+        "management_group_name" : "mgmt-group-uuid",
         "first_name": ppoc.first_name,
         "last_name": ppoc.last_name,
         "country_code": "US",