Merge branch 'staging' into fix-max-width

atst/domain/csp/cloud/azure_cloud_provider.py

@@ -24,6 +24,7 @@ from .models import (
     BillingProfileTenantAccessCSPResult,
     BillingProfileVerificationCSPPayload,
     BillingProfileVerificationCSPResult,
+    CostManagementQueryCSPResult,
     KeyVaultCredentials,
     ManagementGroupCSPResponse,
     ProductPurchaseCSPPayload,
@@ -32,6 +33,7 @@ from .models import (
     ProductPurchaseVerificationCSPResult,
     PrincipalAdminRoleCSPPayload,
     PrincipalAdminRoleCSPResult,
+    ReportingCSPPayload,
     TaskOrderBillingCreationCSPPayload,
     TaskOrderBillingCreationCSPResult,
     TaskOrderBillingVerificationCSPPayload,
@@ -1070,3 +1072,41 @@ class AzureCloudProvider(CloudProviderInterface):
         hashed = sha256_hex(tenant_id)
         raw_creds = self.get_secret(hashed)
         return KeyVaultCredentials(**json.loads(raw_creds))
+
+    def get_reporting_data(self, payload: ReportingCSPPayload):
+        """
+        Queries the Cost Management API for an invoice section's raw reporting data
+        
+        We query at the invoiceSection scope. The full scope path is passed in
+        with the payload at the `invoice_section_id` key.
+        """
+        creds = self._source_tenant_creds(payload.tenant_id)
+        token = self._get_sp_token(
+            payload.tenant_id, creds.tenant_sp_client_id, creds.tenant_sp_key
+        )
+
+        if not token:
+            raise AuthenticationException("Could not retrieve tenant access token")
+
+        headers = {"Authorization": f"Bearer {token}"}
+
+        request_body = {
+            "type": "Usage",
+            "timeframe": "Custom",
+            "timePeriod": {"from": payload.from_date, "to": payload.to_date,},
+            "dataset": {
+                "granularity": "Daily",
+                "aggregation": {"totalCost": {"name": "PreTaxCost", "function": "Sum"}},
+                "grouping": [{"type": "Dimension", "name": "InvoiceId"}],
+            },
+        }
+        cost_mgmt_url = (
+            f"/providers/Microsoft.CostManagement/query?api-version=2019-11-01"
+        )
+        result = self.sdk.requests.post(
+            f"{self.sdk.cloud.endpoints.resource_manager}{payload.invoice_section_id}{cost_mgmt_url}",
+            json=request_body,
+            headers=headers,
+        )
+        if result.ok:
+            return CostManagementQueryCSPResult(**result.json())

atst/domain/csp/cloud/mock_cloud_provider.py

@@ -25,12 +25,15 @@ from .models import (
     BillingProfileTenantAccessCSPResult,
     BillingProfileVerificationCSPPayload,
     BillingProfileVerificationCSPResult,
+    CostManagementQueryCSPResult,
+    CostManagementQueryProperties,
     ProductPurchaseCSPPayload,
     ProductPurchaseCSPResult,
     ProductPurchaseVerificationCSPPayload,
     ProductPurchaseVerificationCSPResult,
     PrincipalAdminRoleCSPPayload,
     PrincipalAdminRoleCSPResult,
+    ReportingCSPPayload,
     SubscriptionCreationCSPPayload,
     SubscriptionCreationCSPResult,
     SubscriptionVerificationCSPPayload,
@@ -487,3 +490,25 @@ class MockCloudProvider(CloudProviderInterface):
 
     def update_tenant_creds(self, tenant_id, secret):
         return secret
+
+    def get_reporting_data(self, payload: ReportingCSPPayload):
+        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)
+        object_id = str(uuid4())
+
+        properties = CostManagementQueryProperties(
+            **dict(
+                columns=[
+                    {"name": "PreTaxCost", "type": "Number"},
+                    {"name": "UsageDate", "type": "Number"},
+                    {"name": "InvoiceId", "type": "String"},
+                    {"name": "Currency", "type": "String"},
+                ],
+                rows=[],
+            )
+        )
+
+        return CostManagementQueryCSPResult(
+            **dict(name=object_id, properties=properties,)
+        )

atst/domain/csp/cloud/models.py

@@ -499,3 +499,34 @@ class UserCSPPayload(BaseCSPPayload):
 
 class UserCSPResult(AliasModel):
     id: str
+
+
+class QueryColumn(AliasModel):
+    name: str
+    type: str
+
+
+class CostManagementQueryProperties(AliasModel):
+    columns: List[QueryColumn]
+    rows: List[Optional[list]]
+
+
+class CostManagementQueryCSPResult(AliasModel):
+    name: str
+    properties: CostManagementQueryProperties
+
+
+class ReportingCSPPayload(BaseCSPPayload):
+    invoice_section_id: str
+    from_date: str
+    to_date: str
+
+    @root_validator(pre=True)
+    def extract_invoice_section(cls, values):
+        try:
+            values["invoice_section_id"] = values["billing_profile_properties"][
+                "invoice_sections"
+            ][0]["invoice_section_id"]
+            return values
+        except (KeyError, IndexError):
+            raise ValueError("Invoice section ID not present in payload")

atst/domain/portfolios/portfolios.py

@@ -15,6 +15,8 @@ from atst.models import (
     Permissions,
     PortfolioRole,
     PortfolioRoleStatus,
+    TaskOrder,
+    CLIN,
 )
 
 from .query import PortfoliosQuery, PortfolioStateMachinesQuery
@@ -144,7 +146,7 @@ class Portfolios(object):
         return db.session.query(Portfolio.id)
 
     @classmethod
-    def get_portfolios_pending_provisioning(cls) -> List[UUID]:
+    def get_portfolios_pending_provisioning(cls, now) -> List[UUID]:
         """
         Any portfolio with a corresponding State Machine that is either:
             not started yet,
@@ -153,22 +155,18 @@ class Portfolios(object):
         """
 
         results = (
-            cls.base_provision_query()
+            db.session.query(Portfolio.id)
             .join(PortfolioStateMachine)
+            .join(TaskOrder)
+            .join(CLIN)
+            .filter(Portfolio.deleted == False)
+            .filter(CLIN.start_date <= now)
+            .filter(CLIN.end_date > now)
             .filter(
                 or_(
                     PortfolioStateMachine.state == FSMStates.UNSTARTED,
-                    PortfolioStateMachine.state == FSMStates.FAILED,
+                    PortfolioStateMachine.state.like("%CREATED"),
-                    PortfolioStateMachine.state == FSMStates.TENANT_FAILED,
                 )
             )
         )
         return [id_ for id_, in results]
-
-        # db.session.query(PortfolioStateMachine).\
-        #        filter(
-        #            or_(
-        #                PortfolioStateMachine.state==FSMStates.UNSTARTED,
-        #                PortfolioStateMachine.state==FSMStates.UNSTARTED,
-        #            )
-        #        ).all()

atst/jobs.py

@@ -203,7 +203,7 @@ def dispatch_provision_portfolio(self):
     """
     Iterate over portfolios with a corresponding State Machine that have not completed.
     """
-    for portfolio_id in Portfolios.get_portfolios_pending_provisioning():
+    for portfolio_id in Portfolios.get_portfolios_pending_provisioning(pendulum.now()):
         provision_portfolio.delay(portfolio_id=portfolio_id)
 
 

atst/models/portfolio_state_machine.py

@@ -175,11 +175,14 @@ class PortfolioStateMachine(
             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)

tests/domain/cloud/test_azure_csp.py

@@ -5,6 +5,8 @@ from uuid import uuid4
 import pytest
 from tests.factories import ApplicationFactory, EnvironmentFactory
 from tests.mock_azure import AUTH_CREDENTIALS, mock_azure
+import pendulum
+import pydantic
 
 from atst.domain.csp.cloud import AzureCloudProvider
 from atst.domain.csp.cloud.models import (
@@ -20,10 +22,12 @@ from atst.domain.csp.cloud.models import (
     BillingProfileTenantAccessCSPResult,
     BillingProfileVerificationCSPPayload,
     BillingProfileVerificationCSPResult,
+    CostManagementQueryCSPResult,
     ProductPurchaseCSPPayload,
     ProductPurchaseCSPResult,
     ProductPurchaseVerificationCSPPayload,
     ProductPurchaseVerificationCSPResult,
+    ReportingCSPPayload,
     SubscriptionCreationCSPPayload,
     SubscriptionCreationCSPResult,
     SubscriptionVerificationCSPPayload,
@@ -718,3 +722,77 @@ def test_create_subscription_verification(mock_azure: AzureCloudProvider):
             payload
         )
         assert result.subscription_id == "60fbbb72-0516-4253-ab18-c92432ba3230"
+
+
+def test_get_reporting_data(mock_azure: AzureCloudProvider):
+    mock_result = Mock()
+    mock_result.json.return_value = {
+        "eTag": None,
+        "id": "providers/Microsoft.Billing/billingAccounts/52865e4c-52e8-5a6c-da6b-c58f0814f06f:7ea5de9d-b8ce-4901-b1c5-d864320c7b03_2019-05-31/billingProfiles/XQDJ-6LB4-BG7-TGB/invoiceSections/P73M-XC7J-PJA-TGB/providers/Microsoft.CostManagement/query/e82d0cda-2ffb-4476-a98a-425c83c216f9",
+        "location": None,
+        "name": "e82d0cda-2ffb-4476-a98a-425c83c216f9",
+        "properties": {
+            "columns": [
+                {"name": "PreTaxCost", "type": "Number"},
+                {"name": "UsageDate", "type": "Number"},
+                {"name": "InvoiceId", "type": "String"},
+                {"name": "Currency", "type": "String"},
+            ],
+            "nextLink": None,
+            "rows": [],
+        },
+        "sku": None,
+        "type": "Microsoft.CostManagement/query",
+    }
+    mock_result.ok = True
+    mock_azure.sdk.requests.post.return_value = mock_result
+    mock_azure = mock_get_secret(mock_azure)
+
+    # Subset of a profile's CSP data that we care about for reporting
+    csp_data = {
+        "tenant_id": "6d2d2d6c-a6d6-41e1-8bb1-73d11475f8f4",
+        "billing_profile_properties": {
+            "invoice_sections": [
+                {
+                    "invoice_section_id": "providers/Microsoft.Billing/billingAccounts/52865e4c-52e8-5a6c-da6b-c58f0814f06f:7ea5de9d-b8ce-4901-b1c5-d864320c7b03_2019-05-31/billingProfiles/XQDJ-6LB4-BG7-TGB/invoiceSections/P73M-XC7J-PJA-TGB",
+                }
+            ],
+        },
+    }
+
+    data: CostManagementQueryCSPResult = mock_azure.get_reporting_data(
+        ReportingCSPPayload(
+            from_date=pendulum.now().subtract(years=1).add(days=1).format("YYYY-MM-DD"),
+            to_date=pendulum.now().format("YYYY-MM-DD"),
+            **csp_data,
+        )
+    )
+
+    assert isinstance(data, CostManagementQueryCSPResult)
+    assert data.name == "e82d0cda-2ffb-4476-a98a-425c83c216f9"
+    assert len(data.properties.columns) == 4
+
+
+def test_get_reporting_data_malformed_payload(mock_azure: AzureCloudProvider):
+    mock_result = Mock()
+    mock_result.ok = True
+    mock_azure.sdk.requests.post.return_value = mock_result
+    mock_azure = mock_get_secret(mock_azure)
+
+    # Malformed csp_data payloads that should throw pydantic validation errors
+    index_error = {
+        "tenant_id": "6d2d2d6c-a6d6-41e1-8bb1-73d11475f8f4",
+        "billing_profile_properties": {"invoice_sections": [],},
+    }
+    key_error = {
+        "tenant_id": "6d2d2d6c-a6d6-41e1-8bb1-73d11475f8f4",
+        "billing_profile_properties": {"invoice_sections": [{}],},
+    }
+
+    for malformed_payload in [key_error, index_error]:
+        with pytest.raises(pydantic.ValidationError):
+            assert mock_azure.get_reporting_data(
+                ReportingCSPPayload(
+                    from_date="foo", to_date="bar", **malformed_payload,
+                )
+            )

tests/domain/test_environments.py

@@ -1,5 +1,4 @@
 import pytest
-import pendulum
 from uuid import uuid4
 
 from atst.domain.environments import Environments
@@ -14,6 +13,7 @@ from tests.factories import (
     EnvironmentRoleFactory,
     ApplicationRoleFactory,
 )
+from tests.utils import EnvQueryTest
 
 
 def test_create_environments():
@@ -119,40 +119,6 @@ def test_update_does_not_duplicate_names_within_application():
         Environments.update(dupe_env, name)
 
 
-class EnvQueryTest:
-    @property
-    def NOW(self):
-        return pendulum.now()
-
-    @property
-    def YESTERDAY(self):
-        return self.NOW.subtract(days=1)
-
-    @property
-    def TOMORROW(self):
-        return self.NOW.add(days=1)
-
-    def create_portfolio_with_clins(self, start_and_end_dates, env_data=None):
-        env_data = env_data or {}
-        return PortfolioFactory.create(
-            applications=[
-                {
-                    "name": "Mos Eisley",
-                    "description": "Where Han shot first",
-                    "environments": [{"name": "thebar", **env_data}],
-                }
-            ],
-            task_orders=[
-                {
-                    "create_clins": [
-                        {"start_date": start_date, "end_date": end_date}
-                        for (start_date, end_date) in start_and_end_dates
-                    ]
-                }
-            ],
-        )
-
-
 class TestGetEnvironmentsPendingCreate(EnvQueryTest):
     def test_with_expired_clins(self, session):
         self.create_portfolio_with_clins([(self.YESTERDAY, self.YESTERDAY)])

tests/domain/test_portfolios.py

@@ -26,6 +26,7 @@ from tests.factories import (
     PortfolioStateMachineFactory,
     get_all_portfolio_permission_sets,
 )
+from tests.utils import EnvQueryTest
 
 
 @pytest.fixture(scope="function")
@@ -263,10 +264,44 @@ def test_create_state_machine(portfolio):
     assert fsm
 
 
-def test_get_portfolios_pending_provisioning(session):
+class TestGetPortfoliosPendingCreate(EnvQueryTest):
-    for x in range(5):
+    def test_finds_unstarted(self):
-        portfolio = PortfolioFactory.create()
+        for x in range(5):
-        sm = PortfolioStateMachineFactory.create(portfolio=portfolio)
+            if x == 2:
-        if x == 2:
+                state = "COMPLETED"
-            sm.state = FSMStates.COMPLETED
+            else:
-    assert len(Portfolios.get_portfolios_pending_provisioning()) == 4
+                state = "UNSTARTED"
+            self.create_portfolio_with_clins(
+                [(self.YESTERDAY, self.TOMORROW)], state_machine_status=state
+            )
+        assert len(Portfolios.get_portfolios_pending_provisioning(self.NOW)) == 4
+
+    def test_finds_created(self):
+        self.create_portfolio_with_clins(
+            [(self.YESTERDAY, self.TOMORROW)], state_machine_status="TENANT_CREATED"
+        )
+        assert len(Portfolios.get_portfolios_pending_provisioning(self.NOW)) == 1
+
+    def test_does_not_find_failed(self):
+        self.create_portfolio_with_clins(
+            [(self.YESTERDAY, self.TOMORROW)], state_machine_status="TENANT_FAILED"
+        )
+        assert len(Portfolios.get_portfolios_pending_provisioning(self.NOW)) == 0
+
+    def test_with_expired_clins(self):
+        self.create_portfolio_with_clins([(self.YESTERDAY, self.YESTERDAY)])
+        assert len(Portfolios.get_portfolios_pending_provisioning(self.NOW)) == 0
+
+    def test_with_active_clins(self):
+        portfolio = self.create_portfolio_with_clins([(self.YESTERDAY, self.TOMORROW)])
+        Portfolios.get_portfolios_pending_provisioning(self.NOW) == [portfolio.id]
+
+    def test_with_future_clins(self):
+        self.create_portfolio_with_clins([(self.TOMORROW, self.TOMORROW)])
+        assert len(Portfolios.get_portfolios_pending_provisioning(self.NOW)) == 0
+
+    def test_with_already_provisioned_env(self):
+        self.create_portfolio_with_clins(
+            [(self.YESTERDAY, self.TOMORROW)], env_data={"cloud_id": uuid4().hex}
+        )
+        assert len(Portfolios.get_portfolios_pending_provisioning(self.NOW)) == 0

tests/test_jobs.py

@@ -286,9 +286,19 @@ def test_create_environment_no_dupes(session, celery_app, celery_worker):
     assert environment.claimed_until == None
 
 
-def test_dispatch_provision_portfolio(
+def test_dispatch_provision_portfolio(csp, monkeypatch):
-    csp, session, portfolio, celery_app, celery_worker, monkeypatch
+    portfolio = PortfolioFactory.create(
-):
+        task_orders=[
+            {
+                "create_clins": [
+                    {
+                        "start_date": pendulum.now().subtract(days=1),
+                        "end_date": pendulum.now().add(days=1),
+                    }
+                ]
+            }
+        ],
+    )
     sm = PortfolioStateMachineFactory.create(portfolio=portfolio)
     mock = Mock()
     monkeypatch.setattr("atst.jobs.provision_portfolio", mock)

tests/utils.py

@@ -5,9 +5,12 @@ from unittest.mock import Mock
 from OpenSSL import crypto
 from cryptography.hazmat.backends import default_backend
 from flask import template_rendered
+import pendulum
 
 from atst.utils.notification_sender import NotificationSender
 
+import tests.factories as factories
+
 
 @contextmanager
 def captured_templates(app):
@@ -62,3 +65,40 @@ def make_crl_list(x509_obj, x509_path):
     issuer = x509_obj.issuer.public_bytes(default_backend())
     filename = os.path.basename(x509_path)
     return [(filename, issuer.hex())]
+
+
+class EnvQueryTest:
+    @property
+    def NOW(self):
+        return pendulum.now()
+
+    @property
+    def YESTERDAY(self):
+        return self.NOW.subtract(days=1)
+
+    @property
+    def TOMORROW(self):
+        return self.NOW.add(days=1)
+
+    def create_portfolio_with_clins(
+        self, start_and_end_dates, env_data=None, state_machine_status=None
+    ):
+        env_data = env_data or {}
+        return factories.PortfolioFactory.create(
+            state=state_machine_status,
+            applications=[
+                {
+                    "name": "Mos Eisley",
+                    "description": "Where Han shot first",
+                    "environments": [{"name": "thebar", **env_data}],
+                }
+            ],
+            task_orders=[
+                {
+                    "create_clins": [
+                        {"start_date": start_date, "end_date": end_date}
+                        for (start_date, end_date) in start_and_end_dates
+                    ]
+                }
+            ],
+        )