@W-18129346 - Update tests

Author: vsbharath

cumulusci/tasks/preflight/tests/test_licenses.py

@@ -13,19 +13,23 @@ class TestLicensePreflights:
     def test_license_preflight(self):
         task = create_task(GetAvailableLicenses, {})
         task._init_api = Mock()
+
+        # Mock license data with usage
         task._init_api.return_value.query.return_value = {
-            "totalSize": 2,
+            "totalSize": 3,
             "records": [
-                {"LicenseDefinitionKey": "TEST1"},
-                {"LicenseDefinitionKey": "TEST2"},
+                {"LicenseDefinitionKey": "TEST1", "TotalLicenses": 100, "UsedLicenses": 50},
+                {"LicenseDefinitionKey": "TEST2", "TotalLicenses": 10, "UsedLicenses": 10},
+                {"LicenseDefinitionKey": "TEST3", "TotalLicenses": 20, "UsedLicenses": 15},
             ],
         }
-        task()
 
+        task()
         task._init_api.return_value.query.assert_called_once_with(
-            "SELECT LicenseDefinitionKey FROM UserLicense"
+            "SELECT Id, LicenseDefinitionKey, TotalLicenses, UsedLicenses FROM UserLicense WHERE Status = 'Active'"
         )
-        assert task.return_values == ["TEST1", "TEST2"]
+        # Only TEST1 and TEST3 have available licenses
+        assert task.return_values == ["TEST1", "TEST3"]
 
     def test_psl_preflight(self):
         task = create_task(GetAvailablePermissionSetLicenses, {})
@@ -47,25 +51,18 @@ def test_psl_preflight(self):
 
     def test_assigned_permsetlicense_preflight(self):
         task = create_task(GetPermissionLicenseSetAssignments, {})
+        task.org_config.user_id = "USER_ID"
         task._init_api = Mock()
+
         task._init_api.return_value.query_all.return_value = {
             "totalSize": 2,
             "done": True,
             "records": [
-                {
-                    "PermissionSetLicense": {
-                        "MasterLabel": "Document Checklist",
-                        "DeveloperName": "DocumentChecklist",
-                    },
-                },
-                {
-                    "PermissionSetLicense": {
-                        "MasterLabel": "Einstein Analytics Plus Admin",
-                        "DeveloperName": "EinsteinAnalyticsPlusAdmin",
-                    },
-                },
+                {"PermissionSetLicense": {"DeveloperName": "DocumentChecklist"}},
+                {"PermissionSetLicense": {"DeveloperName": "EinsteinAnalyticsPlusAdmin"}},
             ],
         }
+
         task()
 
         task._init_api.return_value.query_all.assert_called_once_with(
@@ -80,16 +77,30 @@ def test_permsets_preflight(self):
         task = create_task(GetAvailablePermissionSets, {})
         task._init_api = Mock()
 
-        task._init_api.return_value.query_all.return_value = {
+        # Mock user licenses
+        task._init_api.return_value.query.return_value = {
             "totalSize": 2,
             "records": [
-                {"Name": "TEST1"},
-                {"Name": "TEST2"},
+                {"Id": "L1", "LicenseDefinitionKey": "Platform", "TotalLicenses": 100, "UsedLicenses": 90},
+                {"Id": "L2", "LicenseDefinitionKey": "Force", "TotalLicenses": 50, "UsedLicenses": 50},  # fully used
             ],
         }
-        task()
 
+        # Mock permission sets
+        task._init_api.return_value.query_all.return_value = {
+            "totalSize": 3,
+            "records": [
+                {"Name": "TEST1", "LicenseId": None},
+                {"Name": "TEST2", "LicenseId": "L1"},  # available
+                {"Name": "TEST3", "LicenseId": "L2"},  # not available
+            ],
+        }
+
+        task()
+        task._init_api.return_value.query.assert_called_once_with(
+            "SELECT Id, LicenseDefinitionKey, TotalLicenses, UsedLicenses FROM UserLicense WHERE Status = 'Active'"
+        )
         task._init_api.return_value.query_all.assert_called_once_with(
-            "SELECT Name FROM PermissionSet"
+            "SELECT LicenseId, Name FROM PermissionSet"
         )
         assert task.return_values == ["TEST1", "TEST2"]