From 466a575229f68c2dd46b00f4d3c6293c2f2c4d72 Mon Sep 17 00:00:00 2001
From: tomdds <tom.chandler-ctr@friends.dds.mil>
Date: Sun, 26 Jan 2020 15:17:53 -0500
Subject: [PATCH] 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                   | 19 ------------
 .../csp/cloud/cloud_provider_interface.py     |  1 +
 atst/models/portfolio_state_machine.py        | 30 ++++++++++++++-----
 tests/domain/test_portfolio_state_machine.py  |  2 +-
 4 files changed, 25 insertions(+), 27 deletions(-)

diff --git a/atst/domain/csp/__init__.py b/atst/domain/csp/__init__.py
index f15ac1cd..d886f8a2 100644
--- a/atst/domain/csp/__init__.py
+++ b/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)
diff --git a/atst/domain/csp/cloud/cloud_provider_interface.py b/atst/domain/csp/cloud/cloud_provider_interface.py
index 7f975c07..087aada0 100644
--- a/atst/domain/csp/cloud/cloud_provider_interface.py
+++ b/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()
+
diff --git a/atst/models/portfolio_state_machine.py b/atst/models/portfolio_state_machine.py
index 1390ceb1..cf42710b 100644
--- a/atst/models/portfolio_state_machine.py
+++ b/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}"
diff --git a/tests/domain/test_portfolio_state_machine.py b/tests/domain/test_portfolio_state_machine.py
index 97056ae7..2e412653 100644
--- a/tests/domain/test_portfolio_state_machine.py
+++ b/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