diff --git a/atst/domain/csp/cloud/azure_cloud_provider.py b/atst/domain/csp/cloud/azure_cloud_provider.py index dc5d92c7..22551ed2 100644 --- a/atst/domain/csp/cloud/azure_cloud_provider.py +++ b/atst/domain/csp/cloud/azure_cloud_provider.py @@ -8,8 +8,8 @@ from atst.utils import sha256_hex from .cloud_provider_interface import CloudProviderInterface from .exceptions import ( AuthenticationException, - UserProvisioningException, SecretException, + UserProvisioningException, ) from .models import ( AdminRoleDefinitionCSPPayload, @@ -24,6 +24,7 @@ from .models import ( BillingProfileTenantAccessCSPResult, BillingProfileVerificationCSPPayload, BillingProfileVerificationCSPResult, + CostManagementQueryCSPResult, EnvironmentCSPPayload, EnvironmentCSPResult, KeyVaultCredentials, @@ -33,6 +34,7 @@ from .models import ( ProductPurchaseCSPResult, ProductPurchaseVerificationCSPPayload, ProductPurchaseVerificationCSPResult, + ReportingCSPPayload, SubscriptionCreationCSPPayload, SubscriptionCreationCSPResult, SubscriptionVerificationCSPPayload, @@ -1083,3 +1085,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()) diff --git a/atst/domain/csp/cloud/mock_cloud_provider.py b/atst/domain/csp/cloud/mock_cloud_provider.py index 4955f327..4e4532d3 100644 --- a/atst/domain/csp/cloud/mock_cloud_provider.py +++ b/atst/domain/csp/cloud/mock_cloud_provider.py @@ -23,12 +23,15 @@ from .models import ( BillingProfileTenantAccessCSPResult, BillingProfileVerificationCSPPayload, BillingProfileVerificationCSPResult, + CostManagementQueryCSPResult, + CostManagementQueryProperties, ProductPurchaseCSPPayload, ProductPurchaseCSPResult, ProductPurchaseVerificationCSPPayload, ProductPurchaseVerificationCSPResult, PrincipalAdminRoleCSPPayload, PrincipalAdminRoleCSPResult, + ReportingCSPPayload, SubscriptionCreationCSPPayload, SubscriptionCreationCSPResult, SubscriptionVerificationCSPPayload, @@ -471,3 +474,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,) + ) diff --git a/atst/domain/csp/cloud/models.py b/atst/domain/csp/cloud/models.py index 05f3c866..358f7934 100644 --- a/atst/domain/csp/cloud/models.py +++ b/atst/domain/csp/cloud/models.py @@ -507,3 +507,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") diff --git a/atst/domain/portfolios/portfolios.py b/atst/domain/portfolios/portfolios.py index 1254ac71..b8663730 100644 --- a/atst/domain/portfolios/portfolios.py +++ b/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 == FSMStates.TENANT_FAILED, + PortfolioStateMachine.state.like("%CREATED"), ) ) ) return [id_ for id_, in results] - - # db.session.query(PortfolioStateMachine).\ - # filter( - # or_( - # PortfolioStateMachine.state==FSMStates.UNSTARTED, - # PortfolioStateMachine.state==FSMStates.UNSTARTED, - # ) - # ).all() diff --git a/atst/jobs.py b/atst/jobs.py index 3453f81b..f7ac2df9 100644 --- a/atst/jobs.py +++ b/atst/jobs.py @@ -171,7 +171,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) diff --git a/atst/models/portfolio_state_machine.py b/atst/models/portfolio_state_machine.py index 14e9c01d..f5c1a461 100644 --- a/atst/models/portfolio_state_machine.py +++ b/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) diff --git a/tests/domain/cloud/test_azure_csp.py b/tests/domain/cloud/test_azure_csp.py index 2c946f2b..c58944b9 100644 --- a/tests/domain/cloud/test_azure_csp.py +++ b/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,6 +22,7 @@ from atst.domain.csp.cloud.models import ( BillingProfileTenantAccessCSPResult, BillingProfileVerificationCSPPayload, BillingProfileVerificationCSPResult, + CostManagementQueryCSPResult, EnvironmentCSPPayload, EnvironmentCSPResult, PrincipalAdminRoleCSPPayload, @@ -28,6 +31,7 @@ from atst.domain.csp.cloud.models import ( ProductPurchaseCSPResult, ProductPurchaseVerificationCSPPayload, ProductPurchaseVerificationCSPResult, + ReportingCSPPayload, SubscriptionCreationCSPPayload, SubscriptionCreationCSPResult, SubscriptionVerificationCSPPayload, @@ -753,3 +757,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, + ) + ) diff --git a/tests/domain/test_environments.py b/tests/domain/test_environments.py index a86c6527..4f1308de 100644 --- a/tests/domain/test_environments.py +++ b/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,44 +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, app_data=None - ): - env_data = env_data or {} - app_data = app_data or {} - return PortfolioFactory.create( - applications=[ - { - "name": "Mos Eisley", - "description": "Where Han shot first", - "environments": [{"name": "thebar", **env_data}], - **app_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)]) diff --git a/tests/domain/test_portfolios.py b/tests/domain/test_portfolios.py index ff8ccacb..1093253b 100644 --- a/tests/domain/test_portfolios.py +++ b/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): - for x in range(5): - portfolio = PortfolioFactory.create() - sm = PortfolioStateMachineFactory.create(portfolio=portfolio) - if x == 2: - sm.state = FSMStates.COMPLETED - assert len(Portfolios.get_portfolios_pending_provisioning()) == 4 +class TestGetPortfoliosPendingCreate(EnvQueryTest): + def test_finds_unstarted(self): + for x in range(5): + if x == 2: + state = "COMPLETED" + else: + 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 diff --git a/tests/test_jobs.py b/tests/test_jobs.py index 8e3d5b1d..a5549407 100644 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -253,9 +253,19 @@ def test_create_environment_no_dupes(session, celery_app, celery_worker): assert environment.claimed_until == None -def test_dispatch_provision_portfolio( - csp, session, portfolio, celery_app, celery_worker, monkeypatch -): +def test_dispatch_provision_portfolio(csp, 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) diff --git a/tests/utils.py b/tests/utils.py index 66bf2b18..ca577cff 100644 --- a/tests/utils.py +++ b/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,46 @@ 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, + app_data=None, + state_machine_status=None, + ): + env_data = env_data or {} + app_data = app_data or {} + return factories.PortfolioFactory.create( + state=state_machine_status, + applications=[ + { + "name": "Mos Eisley", + "description": "Where Han shot first", + "environments": [{"name": "thebar", **env_data}], + **app_data, + } + ], + task_orders=[ + { + "create_clins": [ + {"start_date": start_date, "end_date": end_date} + for (start_date, end_date) in start_and_end_dates + ] + } + ], + )