Used MockInstallation consistently instead of create_autospec(Installation)

tests/unit/aws/test_access.py

@@ -4,26 +4,17 @@
 from unittest.mock import MagicMock, call, create_autospec
 
 import pytest
-from databricks.labs.blueprint.installation import Installation, MockInstallation
+from databricks.labs.blueprint.installation import MockInstallation
 from databricks.labs.blueprint.tui import MockPrompts
 from databricks.labs.lsql.backends import MockBackend
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import ResourceDoesNotExist
 from databricks.sdk.service import iam
-from databricks.sdk.service.catalog import (
-    AwsIamRoleResponse,
-    ExternalLocationInfo,
-    StorageCredentialInfo,
-)
+from databricks.sdk.service.catalog import AwsIamRoleResponse, ExternalLocationInfo, StorageCredentialInfo
 from databricks.sdk.service.compute import InstanceProfile, Policy
-from databricks.sdk.service.sql import GetWorkspaceWarehouseConfigResponse, EndpointConfPair
-
-from databricks.labs.ucx.assessment.aws import (
-    AWSPolicyAction,
-    AWSResources,
-    AWSRole,
-    AWSRoleAction,
-)
+from databricks.sdk.service.sql import EndpointConfPair, GetWorkspaceWarehouseConfigResponse
+
+from databricks.labs.ucx.assessment.aws import AWSPolicyAction, AWSResources, AWSRole, AWSRoleAction
 from databricks.labs.ucx.aws.access import AWSResourcePermissions
 from databricks.labs.ucx.aws.credentials import IamRoleCreation
 from databricks.labs.ucx.aws.locations import AWSExternalLocationsMigration
@@ -146,11 +137,24 @@ def test_create_external_locations(mock_ws, installation_multiple_roles, backend
 
 
 def test_create_external_locations_skip_existing(mock_ws, backend, locations):
-    install = create_autospec(Installation)
-    install.load.return_value = [
-        AWSRoleAction("arn:aws:iam::12345:role/uc-role1", "s3", "WRITE_FILES", "s3://BUCKET1"),
-        AWSRoleAction("arn:aws:iam::12345:role/uc-rolex", "s3", "WRITE_FILES", "s3://BUCKETX"),
-    ]
+    install = MockInstallation(
+        {
+            "uc_roles_access.csv": [
+                {
+                    'privilege': 'WRITE_FILES',
+                    'resource_path': 's3://BUCKET1',
+                    'resource_type': 's3',
+                    'role_arn': 'arn:aws:iam::12345:role/uc-role1',
+                },
+                {
+                    'privilege': 'WRITE_FILES',
+                    'resource_path': 's3://BUCKETX',
+                    'resource_type': 's3',
+                    'role_arn': 'arn:aws:iam::12345:role/uc-role1',
+                },
+            ]
+        }
+    )
     mock_ws.storage_credentials.list.return_value = [
         StorageCredentialInfo(
             id="1",

tests/unit/test_cli.py

@@ -5,11 +5,10 @@
 
 import pytest
 import yaml
-from databricks.labs.blueprint.installation import Installation
 from databricks.labs.blueprint.tui import MockPrompts
 from databricks.sdk import AccountClient, WorkspaceClient
 from databricks.sdk.errors import NotFound
-from databricks.sdk.service import iam, sql, jobs
+from databricks.sdk.service import iam, jobs, sql
 from databricks.sdk.service.catalog import ExternalLocationInfo
 from databricks.sdk.service.compute import ClusterDetails, ClusterSource
 from databricks.sdk.service.workspace import ObjectInfo, ObjectType
@@ -20,32 +19,32 @@
 from databricks.labs.ucx.azure.resources import AzureResources
 from databricks.labs.ucx.cli import (
     alias,
+    assign_metastore,
     cluster_remap,
     create_account_groups,
     create_catalogs_schemas,
+    create_missing_principals,
     create_table_mapping,
     create_uber_principal,
     ensure_assessment_run,
     installations,
+    logs,
     manual_workspace_info,
     migrate_credentials,
     migrate_locations,
+    migrate_tables,
     move,
     open_remote_config,
     principal_prefix_access,
     repair_run,
     revert_cluster_remap,
     revert_migrated_tables,
+    show_all_metastores,
     skip,
     sync_workspace_info,
     validate_external_locations,
     validate_groups_membership,
     workflows,
-    logs,
-    show_all_metastores,
-    assign_metastore,
-    migrate_tables,
-    create_missing_principals,
 )
 from databricks.labs.ucx.contexts.account_cli import AccountContext
 from databricks.labs.ucx.contexts.workspace_cli import WorkspaceContext
@@ -412,8 +411,6 @@ def test_cluster_remap(ws, caplog):
         ClusterDetails(cluster_id="123", cluster_name="test_cluster", cluster_source=ClusterSource.UI),
         ClusterDetails(cluster_id="1234", cluster_name="test_cluster1", cluster_source=ClusterSource.JOB),
     ]
-    installation = create_autospec(Installation)
-    installation.save.return_value = "a/b/c"
     cluster_remap(ws, prompts)
     assert "Remapping the Clusters to UC" in caplog.messages
 
@@ -422,8 +419,6 @@ def test_cluster_remap_error(ws, caplog):
     prompts = MockPrompts({"Please provide the cluster id's as comma separated value from the above list.*": "1"})
     ws = create_autospec(WorkspaceClient)
     ws.clusters.list.return_value = []
-    installation = create_autospec(Installation)
-    installation.save.return_value = "a/b/c"
     cluster_remap(ws, prompts)
     assert "No cluster information present in the workspace" in caplog.messages
 

tests/unit/test_install.py

@@ -6,16 +6,12 @@
 import pytest
 import yaml
 from databricks.labs.blueprint.installation import Installation, MockInstallation
-from databricks.labs.blueprint.installer import InstallState, RawState
+from databricks.labs.blueprint.installer import InstallState
 from databricks.labs.blueprint.parallel import ManyError
 from databricks.labs.blueprint.tui import MockPrompts
-from databricks.labs.blueprint.wheels import (
-    ProductInfo,
-    WheelsV2,
-    find_project_root,
-)
+from databricks.labs.blueprint.wheels import ProductInfo, WheelsV2, find_project_root
 from databricks.labs.lsql.backends import MockBackend
-from databricks.sdk import WorkspaceClient, AccountClient
+from databricks.sdk import AccountClient, WorkspaceClient
 from databricks.sdk.errors import (  # pylint: disable=redefined-builtin
     AlreadyExists,
     InvalidParameterValue,
@@ -27,19 +23,8 @@
 )
 from databricks.sdk.errors.platform import BadRequest
 from databricks.sdk.service import iam, jobs, sql
-from databricks.sdk.service.compute import (
-    ClusterDetails,
-    CreatePolicyResponse,
-    DataSecurityMode,
-    Policy,
-    State,
-)
-from databricks.sdk.service.jobs import (
-    BaseRun,
-    RunLifeCycleState,
-    RunResultState,
-    RunState,
-)
+from databricks.sdk.service.compute import ClusterDetails, CreatePolicyResponse, DataSecurityMode, Policy, State
+from databricks.sdk.service.jobs import BaseRun, RunLifeCycleState, RunResultState, RunState
 from databricks.sdk.service.provisioning import Workspace
 from databricks.sdk.service.sql import (
     Dashboard,
@@ -56,16 +41,8 @@
 import databricks.labs.ucx.uninstall  # noqa
 from databricks.labs.ucx.config import WorkspaceConfig
 from databricks.labs.ucx.framework.dashboards import DashboardFromFiles
-from databricks.labs.ucx.install import (
-    WorkspaceInstallation,
-    WorkspaceInstaller,
-    extract_major_minor,
-    AccountInstaller,
-)
-from databricks.labs.ucx.installer.workflows import (
-    DeployedWorkflows,
-    WorkflowsDeployment,
-)
+from databricks.labs.ucx.install import AccountInstaller, WorkspaceInstallation, WorkspaceInstaller, extract_major_minor
+from databricks.labs.ucx.installer.workflows import DeployedWorkflows, WorkflowsDeployment
 from databricks.labs.ucx.runtime import Workflows
 
 PRODUCT_INFO = ProductInfo.from_class(WorkspaceConfig)
@@ -645,7 +622,7 @@ def test_remove_database(ws):
             r'Do you want to delete the inventory database.*': 'yes',
         }
     )
-    installation = create_autospec(Installation)
+    installation = MockInstallation()
     config = WorkspaceConfig(inventory_database='ucx')
     workflow_installer = create_autospec(WorkflowsDeployment)
     workspace_installation = WorkspaceInstallation(
@@ -664,7 +641,7 @@ def test_remove_database(ws):
     assert sql_backend.queries == ['DROP SCHEMA IF EXISTS hive_metastore.ucx CASCADE']
     ws.jobs.delete.assert_not_called()
     ws.cluster_policies.delete.assert_called_once()
-    installation.remove.assert_called_once()
+    installation.assert_removed()
     workflow_installer.create_jobs.assert_not_called()
 
 
@@ -677,7 +654,7 @@ def test_remove_jobs_no_state(ws):
             'Do you want to delete the inventory database ucx too?': 'no',
         }
     )
-    installation = create_autospec(Installation)
+    installation = MockInstallation()
     config = WorkspaceConfig(inventory_database='ucx')
     install_state = InstallState.from_installation(installation)
     wheels = create_autospec(WheelsV2)
@@ -698,7 +675,7 @@ def test_remove_jobs_no_state(ws):
     workspace_installation.uninstall()
 
     ws.jobs.delete.assert_not_called()
-    installation.remove.assert_called_once()
+    installation.assert_removed()
     wheels.upload_to_wsfs.assert_not_called()
 
 
@@ -755,7 +732,7 @@ def test_remove_warehouse(ws):
             'Do you want to delete the inventory database ucx too?': 'no',
         }
     )
-    installation = create_autospec(Installation)
+    installation = MockInstallation()
     config = WorkspaceConfig(inventory_database='ucx', warehouse_id="123")
     workflows_installer = create_autospec(WorkflowsDeployment)
     workspace_installation = WorkspaceInstallation(
@@ -772,7 +749,7 @@ def test_remove_warehouse(ws):
     workspace_installation.uninstall()
 
     ws.warehouses.delete.assert_called_once()
-    installation.remove.assert_called_once()
+    installation.assert_removed()
     workflows_installer.create_jobs.assert_not_called()
 
 
@@ -786,7 +763,7 @@ def test_not_remove_warehouse_with_a_different_prefix(ws):
             'Do you want to delete the inventory database ucx too?': 'no',
         }
     )
-    installation = create_autospec(Installation)
+    installation = MockInstallation()
     config = WorkspaceConfig(inventory_database='ucx', warehouse_id="123")
     workflows_installer = create_autospec(WorkflowsDeployment)
     workspace_installation = WorkspaceInstallation(
@@ -804,7 +781,7 @@ def test_not_remove_warehouse_with_a_different_prefix(ws):
 
     ws.warehouses.delete.assert_not_called()
     workflows_installer.create_jobs.assert_not_called()
-    installation.remove.assert_called_once()
+    installation.assert_removed()
 
 
 def test_remove_secret_scope(ws, caplog):
@@ -869,7 +846,7 @@ def test_remove_cluster_policy_not_exists(ws, caplog):
             'Do you want to delete the inventory database ucx too?': 'no',
         }
     )
-    installation = create_autospec(Installation)
+    installation = MockInstallation()
     config = WorkspaceConfig(inventory_database='ucx')
     ws.cluster_policies.delete.side_effect = NotFound()
     workflows_installer = create_autospec(WorkflowsDeployment)
@@ -888,7 +865,7 @@ def test_remove_cluster_policy_not_exists(ws, caplog):
         workspace_installation.uninstall()
         assert 'UCX Policy already deleted' in caplog.messages
 
-    installation.remove.assert_called_once()
+    installation.assert_removed()
     workflows_installer.create_jobs.assert_not_called()
 
 
@@ -902,7 +879,7 @@ def test_remove_warehouse_not_exists(ws, caplog):
             'Do you want to delete the inventory database ucx too?': 'no',
         }
     )
-    installation = create_autospec(Installation)
+    installation = MockInstallation()
     config = WorkspaceConfig(inventory_database='ucx')
     workflows_installer = create_autospec(WorkflowsDeployment)
     workspace_installation = WorkspaceInstallation(
@@ -920,7 +897,7 @@ def test_remove_warehouse_not_exists(ws, caplog):
         workspace_installation.uninstall()
         assert 'Error accessing warehouse details' in caplog.messages
 
-    installation.remove.assert_called_once()
+    installation.assert_removed()
     workflows_installer.create_jobs.assert_not_called()
 
 
@@ -1643,9 +1620,6 @@ def test_check_inventory_database_exists(ws, mock_installation):
         }
     )
 
-    installation_type_mock = create_autospec(Installation)
-    installation_type_mock.load.side_effect = NotFound
-
     installation = Installation(ws, 'ucx')
     install = WorkspaceInstaller(ws).replace(
         prompts=prompts,
@@ -1686,8 +1660,15 @@ def test_user_not_admin(ws, mock_installation):
     ],
 )
 def test_validate_step(ws, result_state, expected):
-    installation = create_autospec(Installation)
-    installation.load.return_value = RawState({'jobs': {'assessment': '123'}})
+    installation = MockInstallation(
+        {
+            'state.json': {
+                'resources': {
+                    'jobs': {"assessment": "123"},
+                }
+            }
+        }
+    )
     install_state = InstallState.from_installation(installation)
     deployed = DeployedWorkflows(ws, install_state, timedelta(seconds=1))
     ws.jobs.list_runs.return_value = [