Move portfolio state machine helpers directly to model file to prevent import issues.

Having `get_stage_csp_class` in the csp module meant that any file that interacted with that import path would throw an error in a REPL. This will allow importing of the Azure and Mock providers for interactive dev.

atst/domain/csp/__init__.py

@@ -1,5 +1,3 @@
-import importlib
-
 from .cloud import MockCloudProvider
 from .file_uploads import AzureUploader, MockUploader
 from .reports import MockReportingProvider
@@ -31,20 +29,3 @@ def make_csp_provider(app, csp=None):
         app.csp = MockCSP(app, test_mode=True)
     else:
         app.csp = MockCSP(app)
-
-
-def _stage_to_classname(stage):
-    return "".join(map(lambda word: word.capitalize(), stage.split("_")))
-
-
-def get_stage_csp_class(stage, class_type):
-    """
-    given a stage name and class_type return the class
-    class_type is either 'payload' or 'result'
-
-    """
-    cls_name = f"{_stage_to_classname(stage)}CSP{class_type.capitalize()}"
-    try:
-        return getattr(importlib.import_module("atst.domain.csp.cloud"), cls_name)
-    except AttributeError:
-        print("could not import CSP Result class <%s>" % cls_name)

atst/domain/csp/cloud/cloud_provider_interface.py

@@ -124,3 +124,4 @@ class CloudProviderInterface:
         exception if an error occurs while creating a subscription.
         """
         raise NotImplementedError()
+

atst/models/portfolio_state_machine.py

@@ -1,3 +1,5 @@
+import importlib
+
 from sqlalchemy import Column, ForeignKey, Enum as SQLAEnum
 from sqlalchemy.orm import relationship, reconstructor
 from sqlalchemy.dialects.postgresql import UUID
@@ -8,8 +10,7 @@ from transitions.extensions.states import add_state_features, Tags
 
 from flask import current_app as app
 
-from atst.domain.csp.cloud import ConnectionException, UnknownServerException
-from atst.domain.csp import MockCSP, AzureCSP, get_stage_csp_class
+from atst.domain.csp.cloud.exceptions import ConnectionException, UnknownServerException
 from atst.database import db
 from atst.models.types import Id
 from atst.models.base import Base
@@ -17,6 +18,25 @@ import atst.models.mixins as mixins
 from atst.models.mixins.state_machines import FSMStates, AzureStages, _build_transitions
 
 
+def _stage_to_classname(stage):
+    return "".join(map(lambda word: word.capitalize(), stage.split("_")))
+
+
+def get_stage_csp_class(stage, class_type):
+    """
+    given a stage name and class_type return the class
+    class_type is either 'payload' or 'result'
+
+    """
+    cls_name = f"{_stage_to_classname(stage)}CSP{class_type.capitalize()}"
+    try:
+        return getattr(
+            importlib.import_module("atst.domain.csp.cloud.models"), cls_name
+        )
+    except AttributeError:
+        print("could not import CSP Result class <%s>" % cls_name)
+
+
 @add_state_features(Tags)
 class StateMachineWithTags(Machine):
     pass
@@ -138,11 +158,7 @@ class PortfolioStateMachine(
             self.fail_stage(stage)
 
         # TODO: Determine best place to do this, maybe @reconstructor
-        csp = event.kwargs.get("csp")
-        if csp is not None:
-            self.csp = AzureCSP(app).cloud
-        else:
-            self.csp = MockCSP(app).cloud
+        self.csp = app.csp.cloud
 
         try:
             func_name = f"create_{stage}"

tests/domain/test_portfolio_state_machine.py

@@ -10,7 +10,7 @@ from tests.factories import (
 from atst.models import FSMStates, PortfolioStateMachine, TaskOrder
 from atst.models.mixins.state_machines import AzureStages, StageStates, compose_state
 from atst.models.portfolio import Portfolio
-from atst.domain.csp import get_stage_csp_class
+from atst.models.portfolio_state_machine import get_stage_csp_class
 
 # TODO: Write failure case tests