Remove hidden bugs from unit tests by enforcing mocks

pyproject.toml

@@ -67,7 +67,7 @@ dependencies = [
     "mypy~=1.9.0",
     "pylint~=3.1.0",
     "pylint-pytest==2.0.0a0",
-    "databricks-labs-pylint~=0.3.0",
+    "databricks-labs-pylint~=0.4.0",
     "pytest~=8.1.0",
     "pytest-cov~=4.1.0",
     "pytest-mock~=3.14.0",

src/databricks/labs/ucx/aws/credentials.py

@@ -113,13 +113,11 @@ class IamRoleMigration:
     def __init__(
         self,
         installation: Installation,
-        ws: WorkspaceClient,
         resource_permissions: AWSResourcePermissions,
         storage_credential_manager: CredentialManager,
     ):
         self._output_file = "aws_iam_role_migration_result.csv"
         self._installation = installation
-        self._ws = ws
         self._resource_permissions = resource_permissions
         self._storage_credential_manager = storage_credential_manager
 

src/databricks/labs/ucx/contexts/cli_command.py

@@ -162,7 +162,6 @@ def aws_resource_permissions(self):
     def iam_role_migration(self):
         return IamRoleMigration(
             self.installation,
-            self.workspace_client,
             self.aws_resource_permissions,
             self.iam_credential_manager,
         )

tests/integration/aws/test_credentials.py

@@ -32,7 +32,7 @@ def inner(credentials: set[str], read_only=False) -> list[CredentialValidationRe
         location = ExternalLocations(ws, sql_backend, "inventory_schema")
         resource_permissions = AWSResourcePermissions(installation, ws, sql_backend, aws, location, "inventory_schema")
 
-        instance_profile_migration = IamRoleMigration(installation, ws, resource_permissions, CredentialManager(ws))
+        instance_profile_migration = IamRoleMigration(installation, resource_permissions, CredentialManager(ws))
 
         return instance_profile_migration.run(
             MockPrompts({"Above IAM roles will be migrated to UC storage credentials *": "Yes"}),

tests/unit/assessment/test_clusters.py

@@ -142,7 +142,7 @@ def test_try_fetch():
 
 def test_no_isolation_clusters():
     ws = workspace_client_mock(cluster_ids=['no-isolation'])
-    sql_backend = create_autospec(SqlBackend)
+    sql_backend = MockBackend()
     crawler = ClustersCrawler(ws, sql_backend, "ucx")
     result_set = list(crawler.snapshot())
     assert len(result_set) == 1
@@ -151,7 +151,7 @@ def test_no_isolation_clusters():
 
 def test_unsupported_clusters():
     ws = workspace_client_mock(cluster_ids=['legacy-passthrough'])
-    sql_backend = create_autospec(SqlBackend)
+    sql_backend = MockBackend()
     crawler = ClustersCrawler(ws, sql_backend, "ucx")
     result_set = list(crawler.snapshot())
     assert len(result_set) == 1
@@ -163,7 +163,7 @@ def test_policy_crawler():
         policy_ids=['single-user-with-spn', 'single-user-with-spn-policyid', 'single-user-with-spn-no-sparkversion'],
     )
 
-    sql_backend = create_autospec(SqlBackend)
+    sql_backend = MockBackend()
     crawler = PoliciesCrawler(ws, sql_backend, "ucx")
     result_set = list(crawler.snapshot())
     failures = json.loads(result_set[0].failures)

tests/unit/aws/test_credentials.py

@@ -26,19 +26,15 @@ def installation():
     return MockInstallation(DEFAULT_CONFIG)
 
 
-@pytest.fixture
-def ws():
-    return create_autospec(WorkspaceClient)
-
-
 def side_effect_create_aws_storage_credential(name, aws_iam_role, comment, read_only):
     return StorageCredentialInfo(
         name=name, aws_iam_role=AwsIamRoleResponse(role_arn=aws_iam_role.role_arn), comment=comment, read_only=read_only
     )
 
 
 @pytest.fixture
-def credential_manager(ws):
+def credential_manager():
+    ws = create_autospec(WorkspaceClient)
     ws.storage_credentials.list.return_value = [
         StorageCredentialInfo(
             aws_iam_role=AwsIamRoleResponse(role_arn="arn:aws:iam::123456789012:role/example-role-name")
@@ -85,7 +81,7 @@ def test_create_storage_credentials(credential_manager):
 
 
 @pytest.fixture
-def instance_profile_migration(ws, installation, credential_manager):
+def instance_profile_migration(installation, credential_manager):
     def generate_instance_profiles(num_instance_profiles: int):
         arp = create_autospec(AWSResourcePermissions)
         arp.load_uc_compatible_roles.return_value = [
@@ -97,13 +93,12 @@ def generate_instance_profiles(num_instance_profiles: int):
             )
             for i in range(num_instance_profiles)
         ]
-
-        return IamRoleMigration(installation, ws, arp, credential_manager)
+        return IamRoleMigration(installation, arp, credential_manager)
 
     return generate_instance_profiles
 
 
-def test_print_action_plan(caplog, ws, instance_profile_migration, credential_manager):
+def test_print_action_plan(caplog, instance_profile_migration, credential_manager):
     caplog.set_level(logging.INFO)
 
     prompts = MockPrompts({"Above IAM roles will be migrated to UC storage credentials*": "Yes"})
@@ -118,20 +113,21 @@ def test_print_action_plan(caplog, ws, instance_profile_migration, credential_ma
     assert False, "Action plan is not logged"
 
 
-def test_migrate_credential_failed_creation(caplog, ws, instance_profile_migration):
+def test_migrate_credential_failed_creation(caplog, instance_profile_migration):
     caplog.set_level(logging.ERROR)
     prompts = MockPrompts(
         {
             "Above IAM roles will be migrated to UC storage credentials*": "Yes",
         }
     )
+    ws = create_autospec(WorkspaceClient)
     ws.storage_credentials.create.return_value = StorageCredentialInfo(aws_iam_role=None)
     ws.storage_credentials.create.side_effect = None
     instance_profile_migration(1).run(prompts)
     assert "Failed to create storage credential for IAM role: arn:aws:iam::123456789012:role/prefix0" in caplog.messages
 
 
-def test_run_without_confirmation(ws, instance_profile_migration):
+def test_run_without_confirmation(instance_profile_migration):
     prompts = MockPrompts(
         {
             "Above IAM roles will be migrated to UC storage credentials*": "No",
@@ -142,18 +138,18 @@ def test_run_without_confirmation(ws, instance_profile_migration):
 
 
 @pytest.mark.parametrize("num_instance_profiles", [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
-def test_run(ws, instance_profile_migration, num_instance_profiles: int):
+def test_run(instance_profile_migration, num_instance_profiles: int):
     prompts = MockPrompts({"Above IAM roles will be migrated to UC storage credentials*": "Yes"})
     migration = instance_profile_migration(num_instance_profiles)
     results = migration.run(prompts)
     assert len(results) == num_instance_profiles
 
 
-def test_run_no_credential_to_migrate(caplog, ws, installation, credential_manager):
+def test_run_no_credential_to_migrate(caplog, installation, credential_manager):
     caplog.set_level(logging.INFO)
     arp = create_autospec(AWSResourcePermissions)
     arp.load_uc_compatible_roles.return_value = []
-    migration = IamRoleMigration(installation, ws, arp, credential_manager)
+    migration = IamRoleMigration(installation, arp, credential_manager)
     migration.run(MockPrompts({}))
     assert "No IAM Role to migrate" in caplog.messages
 

tests/unit/azure/test_access.py

@@ -32,11 +32,17 @@ def test_save_spn_permissions_no_external_table(caplog):
     location = ExternalLocations(w, backend, "ucx")
     installation = MockInstallation()
     azure_resources = create_autospec(AzureResources)
-    azure_resource_permission = AzureResourcePermissions(installation, w, azure_resources, location)
     azure_resources.storage_accounts.return_value = []
+    azure_resource_permission = AzureResourcePermissions(installation, w, azure_resources, location)
     azure_resource_permission.save_spn_permissions()
     msg = "There are no external table present with azure storage account. Please check if assessment job is run"
     assert [rec.message for rec in caplog.records if msg in rec.message]
+    w.cluster_policies.get.assert_not_called()
+    w.secrets.get_secret.assert_not_called()
+    w.secrets.create_scope.assert_not_called()
+    w.secrets.put_secret.assert_not_called()
+    w.cluster_policies.edit.assert_not_called()
+    w.get_workspace_id.assert_not_called()
 
 
 def test_save_spn_permissions_no_external_tables():
@@ -49,6 +55,12 @@ def test_save_spn_permissions_no_external_tables():
     azure_resource_permission = AzureResourcePermissions(installation, w, azure_resources, location)
     azure_resources.storage_accounts.return_value = []
     assert not azure_resource_permission.save_spn_permissions()
+    w.cluster_policies.get.assert_not_called()
+    w.secrets.get_secret.assert_not_called()
+    w.secrets.create_scope.assert_not_called()
+    w.secrets.put_secret.assert_not_called()
+    w.cluster_policies.edit.assert_not_called()
+    w.get_workspace_id.assert_not_called()
 
 
 def test_save_spn_permissions_no_azure_storage_account():
@@ -63,6 +75,12 @@ def test_save_spn_permissions_no_azure_storage_account():
     azure_resource_permission = AzureResourcePermissions(installation, w, azure_resources, location)
     azure_resources.storage_accounts.return_value = []
     assert not azure_resource_permission.save_spn_permissions()
+    w.cluster_policies.get.assert_not_called()
+    w.secrets.get_secret.assert_not_called()
+    w.secrets.create_scope.assert_not_called()
+    w.secrets.put_secret.assert_not_called()
+    w.cluster_policies.edit.assert_not_called()
+    w.get_workspace_id.assert_not_called()
 
 
 def test_save_spn_permissions_valid_azure_storage_account():
@@ -111,6 +129,12 @@ def test_save_spn_permissions_valid_azure_storage_account():
     ]
     azure_resource_permission = AzureResourcePermissions(installation, w, azure_resources, location)
     azure_resource_permission.save_spn_permissions()
+    w.cluster_policies.get.assert_not_called()
+    w.secrets.get_secret.assert_not_called()
+    w.secrets.create_scope.assert_not_called()
+    w.secrets.put_secret.assert_not_called()
+    w.cluster_policies.edit.assert_not_called()
+    w.get_workspace_id.assert_not_called()
     installation.assert_file_written(
         'azure_storage_account_info.csv',
         [
@@ -153,6 +177,16 @@ def test_create_global_spn_no_policy():
     prompts = MockPrompts({"Enter a name for the uber service principal to be created*": "UCXServicePrincipal"})
     with pytest.raises(ValueError):
         azure_resource_permission.create_uber_principal(prompts)
+    azure_resources.storage_accounts.assert_not_called()
+    azure_resources.create_or_update_access_connector.assert_not_called()
+    azure_resources.role_assignments.assert_not_called()
+    azure_resources.containers.assert_not_called()
+    w.cluster_policies.get.assert_not_called()
+    w.secrets.get_secret.assert_not_called()
+    w.secrets.create_scope.assert_not_called()
+    w.secrets.put_secret.assert_not_called()
+    w.cluster_policies.edit.assert_not_called()
+    w.get_workspace_id.assert_called_once()
 
 
 def test_create_global_spn_spn_present():
@@ -175,6 +209,16 @@ def test_create_global_spn_spn_present():
     prompts = MockPrompts({"Enter a name for the uber service principal to be created*": "UCXServicePrincipal"})
     azure_resource_permission = AzureResourcePermissions(installation, w, azure_resources, location)
     assert not azure_resource_permission.create_uber_principal(prompts)
+    azure_resources.storage_accounts.assert_not_called()
+    azure_resources.create_or_update_access_connector.assert_not_called()
+    azure_resources.role_assignments.assert_not_called()
+    azure_resources.containers.assert_not_called()
+    w.cluster_policies.get.assert_not_called()
+    w.secrets.get_secret.assert_not_called()
+    w.secrets.create_scope.assert_not_called()
+    w.secrets.put_secret.assert_not_called()
+    w.cluster_policies.edit.assert_not_called()
+    w.get_workspace_id.assert_called_once()
 
 
 def test_create_global_spn_no_storage():
@@ -198,6 +242,16 @@ def test_create_global_spn_no_storage():
     azure_resources = create_autospec(AzureResources)
     azure_resource_permission = AzureResourcePermissions(installation, w, azure_resources, location)
     assert not azure_resource_permission.create_uber_principal(prompts)
+    azure_resources.storage_accounts.assert_not_called()
+    azure_resources.create_or_update_access_connector.assert_not_called()
+    azure_resources.role_assignments.assert_not_called()
+    azure_resources.containers.assert_not_called()
+    w.cluster_policies.get.assert_not_called()
+    w.secrets.get_secret.assert_not_called()
+    w.secrets.create_scope.assert_not_called()
+    w.secrets.put_secret.assert_not_called()
+    w.cluster_policies.edit.assert_not_called()
+    w.get_workspace_id.assert_called_once()
 
 
 def test_create_global_spn_cluster_policy_not_found():
@@ -224,6 +278,12 @@ def test_create_global_spn_cluster_policy_not_found():
     azure_resource_permission = AzureResourcePermissions(installation, w, azure_resources, location)
     with pytest.raises(NotFound):
         azure_resource_permission.create_uber_principal(prompts)
+    w.cluster_policies.get.assert_called_once()
+    w.secrets.get_secret.assert_not_called()
+    w.secrets.create_scope.assert_called_with("ucx")
+    w.secrets.put_secret.assert_called_with('ucx', 'uber_principal_secret', string_value='mypwd')
+    w.cluster_policies.edit.assert_not_called()
+    w.get_workspace_id.assert_called_once()
 
 
 def test_create_global_spn():
@@ -302,6 +362,10 @@ def test_create_access_connectors_for_storage_accounts_logs_no_storage_accounts(
     azure_resource_permission = AzureResourcePermissions(installation, w, azure_resources, location)
 
     azure_resource_permission.create_access_connectors_for_storage_accounts()
+
+    w.cluster_policies.get.assert_not_called()
+    w.secrets.get_secret.assert_not_called()
+    w.secrets.create_scope.assert_not_called()
     assert (
         "There are no external table present with azure storage account. Please check if assessment job is run"
         in caplog.messages
@@ -345,6 +409,11 @@ def test_create_access_connectors_for_storage_accounts_one_access_connector():
     azure_resource_permission = AzureResourcePermissions(installation, w, azure_resources, location)
 
     access_connectors = azure_resource_permission.create_access_connectors_for_storage_accounts()
+
+    w.cluster_policies.get.assert_not_called()
+    w.secrets.get_secret.assert_not_called()
+    w.secrets.create_scope.assert_not_called()
+
     assert len(access_connectors) == 1
     assert access_connectors[0].name == "ac-test"
 
@@ -388,3 +457,7 @@ def test_create_access_connectors_for_storage_accounts_log_permission_applied(ca
     with caplog.at_level(logging.DEBUG, logger="databricks.labs.ucx"):
         azure_resource_permission.create_access_connectors_for_storage_accounts()
         assert any("STORAGE_BLOB_DATA_CONTRIBUTOR" in message for message in caplog.messages)
+
+    w.cluster_policies.get.assert_not_called()
+    w.secrets.get_secret.assert_not_called()
+    w.secrets.create_scope.assert_not_called()