Sample create and validate billing profile integration

Adds 2 methods to the azure csp interface to create and subsequently validate creation of the billing profile.

atst/domain/csp/cloud.py

@@ -1,8 +1,8 @@
 import re
-from typing import Dict
+from typing import Dict, List, Optional
 from uuid import uuid4
 
-from pydantic import BaseModel
+from pydantic import BaseModel, validator
 
 from atst.models.user import User
 from atst.models.environment import Environment
@@ -193,6 +193,12 @@ class TenantCSPResult(AliasModel):
 
 
 class BillingProfileAddress(AliasModel):
+    company_name: str
+    address_line_1: str
+    city: str
+    region: str
+    country: str
+    postal_code: str
 
 
 class BillingProfileCLINBudget(AliasModel):
@@ -207,44 +213,60 @@ class BillingProfileCLINBudget(AliasModel):
     """
 
 
-class BillingProfileCSPPayload(
+class BillingProfileCSPPayload(BaseCSPPayload):
-    BaseCSPPayload, BillingProfileAddress, BillingProfileCLINBudget
+    tenant_id: str
-):
+    display_name: str
-    displayName: str
+    enabled_azure_plans: Optional[List[str]]
-    poNumber: str
+    address: BillingProfileAddress
-    invoiceEmailOptIn: str
 
-    """
+    @validator("enabled_azure_plans", pre=True, always=True)
-    {
+    def default_enabled_azure_plans(cls, v):
-        "displayName": "string",
+        """
-        "poNumber": "string",
+        Normally you'd implement this by setting the field with a value of:
-        "address": {
+            dataclasses.field(default_factory=list)
-            "firstName": "string",
+        but that prevents the object from being correctly pickled, so instead we need
-            "lastName": "string",
+        to rely on a validator to ensure this has an empty value when not specified
-            "companyName": "string",
+        """
-            "addressLine1": "string",
+        return v or []
-            "addressLine2": "string",
+
-            "addressLine3": "string",
+
-            "city": "string",
+class BillingProfileCreateCSPResult(AliasModel):
-            "region": "string",
+    location: str
-            "country": "string",
+    retry_after: int
-            "postalCode": "string"
+
-        },
+    class Config:
-        "invoiceEmailOptIn": true,
+        fields = {"location": "Location", "retry_after": "Retry-After"}
-        Note: These last 2 are also the body for adding/updating new TOs/clins
+
-        "enabledAzurePlans": [
+
-            {
+class BillingProfileVerifyCSPPayload(BaseCSPPayload):
-            "skuId": "string"
+    location: str
-            }
+
-        ],
+
-        "clinBudget": {
+class BillingInvoiceSection(AliasModel):
-            "amount": 0,
+    invoice_section_id: str
-            "startDate": "2019-12-18T16:47:40.909Z",
+    invoice_section_name: str
-            "endDate": "2019-12-18T16:47:40.909Z",
+
-            "externalReferenceId": "string"
+    class Config:
+        fields = {"invoice_section_id": "id", "invoice_section_name": "name"}
+
+
+class BillingProfileProperties(AliasModel):
+    address: BillingProfileAddress
+    display_name: str
+    invoice_sections: List[BillingInvoiceSection]
+
+
+class BillingProfileCSPResult(AliasModel):
+    billing_profile_id: str
+    billing_profile_name: str
+    billing_profile_properties: BillingProfileProperties
+
+    class Config:
+        fields = {
+            "billing_profile_id": "id",
+            "billing_profile_name": "name",
+            "billing_profile_properties": "properties",
         }
-    }
-    """
 
 
 class CloudProviderInterface:
@@ -682,8 +704,6 @@ class AzureCloudProvider(CloudProviderInterface):
 
         create_tenant_body = payload.dict(by_alias=True)
 
-        print(create_tenant_body)
-
         create_tenant_headers = {
             "Content-Type": "application/json",
             "Authorization": f"Bearer {sp_token}",
@@ -700,6 +720,53 @@ class AzureCloudProvider(CloudProviderInterface):
         else:
             return self._error(result.json())
 
+    def create_billing_profile(self, payload: BillingProfileCSPPayload):
+        sp_token = self._get_sp_token(payload.creds)
+        if sp_token is None:
+            raise AuthenticationException(
+                "Could not resolve token for billing profile creation"
+            )
+
+        create_billing_account_body = payload.dict(by_alias=True)
+
+        create_billing_account_headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {sp_token}",
+        }
+
+        # TODO: unsure if this is a static value or needs to be constructed/configurable
+        BILLING_ACCOUT_NAME = "7c89b735-b22b-55c0-ab5a-c624843e8bf6:de4416ce-acc6-44b1-8122-c87c4e903c91_2019-05-31"
+        billing_account_create_url = f"https://management.azure.com/providers/Microsoft.Billing/billingAccounts/{BILLING_ACCOUT_NAME}/billingProfiles?api-version=2019-10-01-preview"
+
+        result = self.sdk.requests.post(
+            billing_account_create_url,
+            json=create_billing_account_body,
+            headers=create_billing_account_headers,
+        )
+
+        if result.status_code == 202:
+            return self._ok(BillingProfileCreateCSPResult(**result.headers))
+        else:
+            return self._error(result.json())
+
+    def validate_billing_profile_created(self, payload: BillingProfileVerifyCSPPayload):
+        sp_token = self._get_sp_token(payload.creds)
+        if sp_token is None:
+            raise AuthenticationException(
+                "Could not resolve token for billing profile validation"
+            )
+
+        auth_header = {
+            "Authorization": f"Bearer {sp_token}",
+        }
+
+        result = self.sdk.requests.get(payload.location, headers=auth_header)
+
+        if result.status_code == 200:
+            return self._ok(BillingProfileCSPResult(**result.json()))
+        else:
+            return self._error(result.json())
+
     def create_billing_owner(self, creds, tenant_admin_details):
         # authenticate as tenant_admin
         # create billing owner identity
@@ -722,45 +789,6 @@ class AzureCloudProvider(CloudProviderInterface):
 
         return self.ok()
 
-    def create_billing_profile(self, creds, tenant_admin_details, billing_owner_id):
-        # call billing profile creation endpoint, specifying owner
-        # Payload:
-        """
-        {
-            "displayName": "string",
-            "poNumber": "string",
-            "address": {
-                "firstName": "string",
-                "lastName": "string",
-                "companyName": "string",
-                "addressLine1": "string",
-                "addressLine2": "string",
-                "addressLine3": "string",
-                "city": "string",
-                "region": "string",
-                "country": "string",
-                "postalCode": "string"
-            },
-            "invoiceEmailOptIn": true,
-            Note: These last 2 are also the body for adding/updating new TOs/clins
-            "enabledAzurePlans": [
-                {
-                "skuId": "string"
-                }
-            ],
-            "clinBudget": {
-                "amount": 0,
-                "startDate": "2019-12-18T16:47:40.909Z",
-                "endDate": "2019-12-18T16:47:40.909Z",
-                "externalReferenceId": "string"
-            }
-        }
-        """
-
-        # response will be mostly the same as the body, but we only really care about the id
-        response = {"id": "string"}
-        return self._ok({"billing_profile_id": response["id"]})
-
     def report_clin(self, creds, clin_id, clin_amount, clin_start, clin_end, clin_to):
         # should consumer be responsible for reporting each clin or
         # should this take a list and manage the sequential reporting?

tests/domain/cloud/test_azure_csp.py

@@ -7,6 +7,11 @@ from atst.domain.csp.cloud import (
     AzureCloudProvider,
     TenantCSPResult,
     TenantCSPPayload,
+    BillingProfileCSPPayload,
+    BillingProfileAddress,
+    BillingProfileCreateCSPResult,
+    BillingProfileVerifyCSPPayload,
+    BillingProfileCSPResult,
 )
 
 from tests.mock_azure import mock_azure, AUTH_CREDENTIALS
@@ -156,3 +161,95 @@ def test_create_tenant(mock_azure: AzureCloudProvider):
     print(result)
     body: TenantCSPResult = result.get("body")
     assert body.tenant_id == "60ff9d34-82bf-4f21-b565-308ef0533435"
+
+
+def test_create_billing_profile(mock_azure: AzureCloudProvider):
+    # mock_azure.sdk.adal.AuthenticationContext.return_value.context.acquire_token_with_client_credentials.return_value = {
+    #     "accessToken": "TOKEN"
+    # }
+
+    # mock_result = Mock()
+    # mock_result.headers = {
+    #     "Location": "http://retry-url",
+    #     "Retry-After": "10",
+    # }
+    # mock_result.status_code = 202
+    # mock_azure.sdk.requests.post.return_value = mock_result
+    payload = BillingProfileCSPPayload(
+        **dict(
+            address=dict(
+                address_line_1="123 S Broad Street, Suite 2400",
+                company_name="Promptworks",
+                city="Philadelphia",
+                region="PA",
+                country="US",
+                postal_code="19109",
+            ),
+            creds={"username": "mock-cloud", "pass": "shh"},
+            tenant_id="60ff9d34-82bf-4f21-b565-308ef0533435",
+            display_name="Test Billing Profile",
+        )
+    )
+    result = mock_azure.create_billing_profile(payload)
+    print(result)
+    body: BillingProfileCreateCSPResult = result.get("body")
+    assert body.retry_after == 10
+
+
+def test_validate_billing_profile_creation(mock_azure: AzureCloudProvider):
+    mock_azure.sdk.adal.AuthenticationContext.return_value.context.acquire_token_with_client_credentials.return_value = {
+        "accessToken": "TOKEN"
+    }
+
+    mock_result = Mock()
+    mock_result.status_code = 200
+    mock_result.json.return_value = {
+        "id": "/providers/Microsoft.Billing/billingAccounts/7c89b735-b22b-55c0-ab5a-c624843e8bf6:de4416ce-acc6-44b1-8122-c87c4e903c91_2019-05-31/billingProfiles/XC36-GRNZ-BG7-TGB",
+        "name": "XC36-GRNZ-BG7-TGB",
+        "properties": {
+            "address": {
+                "addressLine1": "123 S Broad Street, Suite 2400",
+                "city": "Philadelphia",
+                "companyName": "Promptworks",
+                "country": "US",
+                "postalCode": "19109",
+                "region": "PA",
+            },
+            "currency": "USD",
+            "displayName": "First Portfolio Billing Profile",
+            "enabledAzurePlans": [],
+            "hasReadAccess": True,
+            "invoiceDay": 5,
+            "invoiceEmailOptIn": False,
+            "invoiceSections": [
+                {
+                    "id": "/providers/Microsoft.Billing/billingAccounts/7c89b735-b22b-55c0-ab5a-c624843e8bf6:de4416ce-acc6-44b1-8122-c87c4e903c91_2019-05-31/billingProfiles/XC36-GRNZ-BG7-TGB/invoiceSections/6HMZ-2HLO-PJA-TGB",
+                    "name": "6HMZ-2HLO-PJA-TGB",
+                    "properties": {"displayName": "First Portfolio Billing Profile"},
+                    "type": "Microsoft.Billing/billingAccounts/billingProfiles/invoiceSections",
+                }
+            ],
+        },
+        "type": "Microsoft.Billing/billingAccounts/billingProfiles",
+    }
+    mock_azure.sdk.requests.get.return_value = mock_result
+
+    payload = BillingProfileVerifyCSPPayload(
+        **dict(
+            creds={
+                "username": "username",
+                "password": "password",
+                "tenant_id": "tenant_id",
+            },
+            location="https://management.azure.com/providers/Microsoft.Billing/billingAccounts/7c89b735-b22b-55c0-ab5a-c624843e8bf6:de4416ce-acc6-44b1-8122-c87c4e903c91_2019-05-31/operationResults/createBillingProfile_478d5706-71f9-4a8b-8d4e-2cbaca27a668?api-version=2019-10-01-preview",
+        )
+    )
+
+    result = mock_azure.validate_billing_profile_created(payload)
+    body: BillingProfileCreateCSPResult = result.get("body")
+    assert body.billing_profile_name == "XC36-GRNZ-BG7-TGB"
+    assert (
+        body.billing_profile_properties.display_name
+        == "First Portfolio Billing Profile"
+    )
+