From 6b8ea4192564b35811b030a0a6f1e365bf9d60eb Mon Sep 17 00:00:00 2001
From: tomdds <tom.chandler-ctr@friends.dds.mil>
Date: Thu, 30 Jan 2020 22:08:48 -0500
Subject: [PATCH] Make AADP purchase quantity configurable

Also remove a few constant params from the payload model.
---
 atst/domain/csp/cloud/azure_cloud_provider.py | 11 +++--------
 atst/domain/csp/cloud/models.py               |  4 ----
 tests/domain/cloud/test_azure_csp.py          |  4 ----
 tests/mock_azure.py                           |  1 +
 4 files changed, 4 insertions(+), 16 deletions(-)

diff --git a/atst/domain/csp/cloud/azure_cloud_provider.py b/atst/domain/csp/cloud/azure_cloud_provider.py
index 3074e221..048b1d5b 100644
--- a/atst/domain/csp/cloud/azure_cloud_provider.py
+++ b/atst/domain/csp/cloud/azure_cloud_provider.py
@@ -99,6 +99,7 @@ class AzureCloudProvider(CloudProviderInterface):
         self.ps_client_id = config["POWERSHELL_CLIENT_ID"]
         self.owner_role_def_id = config["AZURE_OWNER_ROLE_DEF_ID"]
         self.graph_resource = config["AZURE_GRAPH_RESOURCE"]
+        self.default_aadp_qty = config["AZURE_AADP_QTY"]
 
         if azure_sdk_provider is None:
             self.sdk = AzureSDKProvider()
@@ -528,17 +529,11 @@ class AzureCloudProvider(CloudProviderInterface):
                 "Could not resolve token for aad premium product purchase"
             )
 
-        payload_as_dict = payload.dict(by_alias=True)
-
         create_product_purchase_body = {
             "type": "AADPremium",
             "sku": "AADP1",
-            "productProperties": {
-                "beneficiaryTenantId": payload_as_dict["productProperties"][
-                    "beneficiaryTenantId"
-                ],
-            },
-            "quantity": payload_as_dict.get("quantity"),
+            "productProperties": {"beneficiaryTenantId": payload.tenant_id,},
+            "quantity": self.default_aadp_qty,
         }
         create_product_purchase_headers = {
             "Authorization": f"Bearer {sp_token}",
diff --git a/atst/domain/csp/cloud/models.py b/atst/domain/csp/cloud/models.py
index aa24d027..4435093c 100644
--- a/atst/domain/csp/cloud/models.py
+++ b/atst/domain/csp/cloud/models.py
@@ -409,10 +409,6 @@ class KeyVaultCredentials(BaseModel):
 
 
 class ProductPurchaseCSPPayload(BaseCSPPayload):
-    type: str
-    sku: str
-    quantity: int
-    productProperties: Dict
     billing_account_name: str
     billing_profile_name: str
 
diff --git a/tests/domain/cloud/test_azure_csp.py b/tests/domain/cloud/test_azure_csp.py
index eef0a0d8..9e0c3cb0 100644
--- a/tests/domain/cloud/test_azure_csp.py
+++ b/tests/domain/cloud/test_azure_csp.py
@@ -463,10 +463,6 @@ def test_create_product_purchase(mock_azure: AzureCloudProvider):
     payload = ProductPurchaseCSPPayload(
         **dict(
             tenant_id="6d2d2d6c-a6d6-41e1-8bb1-73d11475f8f4",
-            type="AADPremium",
-            sku="AADP1",
-            productProperties={"beneficiaryTenantId": str(uuid4()),},
-            quantity=4,
             billing_account_name="7c89b735-b22b-55c0-ab5a-c624843e8bf6:de4416ce-acc6-44b1-8122-c87c4e903c91_2019-05-31",
             billing_profile_name="KQWI-W2SU-BG7-TGB",
         )
diff --git a/tests/mock_azure.py b/tests/mock_azure.py
index 438ae855..ce85a396 100644
--- a/tests/mock_azure.py
+++ b/tests/mock_azure.py
@@ -12,6 +12,7 @@ AZURE_CONFIG = {
     "POWERSHELL_CLIENT_ID": "MOCK",
     "AZURE_OWNER_ROLE_DEF_ID": "MOCK",
     "AZURE_GRAPH_RESOURCE": "MOCK",
+    "AZURE_AADP_QTY": 5,
 }
 
 AUTH_CREDENTIALS = {