@W-18129346 - Update tests

Author: vsbharath

cumulusci/tasks/preflight/licenses.py

@@ -4,11 +4,7 @@
 class BaseUserLicenseAwareTask(BaseSalesforceApiTask):
     def get_available_user_licenses(self):
         """Fetch active user licenses with availability."""
-        query = """
-            SELECT Id, LicenseDefinitionKey, TotalLicenses, UsedLicenses 
-            FROM UserLicense 
-            WHERE Status = 'Active'
-        """
+        query = "SELECT Id, LicenseDefinitionKey, TotalLicenses, UsedLicenses FROM UserLicense WHERE Status = 'Active'"
         return {
             lic["Id"]: lic
             for lic in self.sf.query(query)["records"]
@@ -18,6 +14,7 @@ def get_available_user_licenses(self):
     def _log_list(self, title, items):
         self.logger.info(f"{title} ({len(items)}):\n" + "\n".join(f"- {item}" for item in items))
 
+
 class GetAvailableLicenses(BaseUserLicenseAwareTask):
     def _run_task(self):
         self.return_values = [

cumulusci/tasks/preflight/tests/test_licenses.py

@@ -12,22 +12,12 @@
 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": 3,
-            "records": [
-                {"LicenseDefinitionKey": "TEST1", "TotalLicenses": 100, "UsedLicenses": 50},
-                {"LicenseDefinitionKey": "TEST2", "TotalLicenses": 10, "UsedLicenses": 10},
-                {"LicenseDefinitionKey": "TEST3", "TotalLicenses": 20, "UsedLicenses": 15},
-            ],
-        }
+        task.get_available_user_licenses = Mock(return_value={
+            "L1": {"LicenseDefinitionKey": "TEST1"},
+            "L3": {"LicenseDefinitionKey": "TEST3"},
+        })
 
         task()
-        task._init_api.return_value.query.assert_called_once_with(
-            "SELECT Id, LicenseDefinitionKey, TotalLicenses, UsedLicenses FROM UserLicense WHERE Status = 'Active'"
-        )
         # Only TEST1 and TEST3 have available licenses
         assert task.return_values == ["TEST1", "TEST3"]
 
@@ -82,31 +72,21 @@ def test_assigned_permsetlicense_preflight(self):
 
     def test_permsets_preflight(self):
         task = create_task(GetAvailablePermissionSets, {})
-        task._init_api = Mock()
+        task.get_available_user_licenses = Mock(return_value={
+            "L1": {"LicenseDefinitionKey": "Platform", "TotalLicenses": 100, "UsedLicenses": 90},
+        })
 
-        # Mock user licenses
-        task._init_api.return_value.query.return_value = {
-            "totalSize": 2,
-            "records": [
-                {"Id": "L1", "LicenseDefinitionKey": "Platform", "TotalLicenses": 100, "UsedLicenses": 90},
-                {"Id": "L2", "LicenseDefinitionKey": "Force", "TotalLicenses": 50, "UsedLicenses": 50},  # fully used
-            ],
-        }
-
-        # Mock permission sets
+        task._init_api = Mock()
         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
+                {"Name": "TEST2", "LicenseId": "L1"},
+                {"Name": "TEST3", "LicenseId": "L2"},  # 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 LicenseId, Name FROM PermissionSet"
         )