Enforce protected-access pylint rule (#956)

Author: nfx

pyproject.toml

@@ -552,7 +552,7 @@ disable = [
     "missing-module-docstring",
     "missing-class-docstring",
     "missing-function-docstring",
-    "protected-access", # TODO: enable back
+#    "protected-access", # TODO: enable back
     "too-few-public-methods",
     "line-too-long",
     "too-many-lines",

src/databricks/labs/ucx/azure/access.py

@@ -80,6 +80,9 @@ def save_spn_permissions(self) -> str | None:
             return None
         return self._installation.save(storage_account_infos, filename=self._filename)
 
+    def load(self):
+        return self._installation.save(list[StoragePermissionMapping], filename=self._filename)
+
     def _get_storage_accounts(self) -> list[str]:
         external_locations = self._locations.snapshot()
         storage_accounts = []

src/databricks/labs/ucx/workspace_access/groups.py

@@ -53,6 +53,9 @@ def partial_info(cls, workspace: iam.Group, account: iam.Group):
             external_id=workspace.external_id,
         )
 
+    def decode_members(self):
+        return [iam.ComplexValue.from_dict(_) for _ in json.loads(self.members)]
+
 
 class MigrationState:
     """Holds migration state of workspace-to-account groups"""
@@ -451,6 +454,14 @@ def validate_group_membership(self) -> list[dict]:
             logger.info("There are groups with different membership between account and workspace")
         return mismatch_group
 
+    def has_workspace_group(self, name):
+        groups = self._workspace_groups_in_workspace()
+        return name in groups
+
+    def has_account_group(self, name):
+        groups = self._account_groups_in_workspace()
+        return name in groups
+
     def _workspace_groups_in_workspace(self) -> dict[str, Group]:
         attributes = "id,displayName,meta,externalId,members,roles,entitlements"
         groups = {}

tests/integration/azure/test_access.py

@@ -22,9 +22,11 @@ def test_azure_storage_accounts(ws, sql_backend, inventory_schema, make_random):
     location = ExternalLocations(ws, sql_backend, inventory_schema)
     installation = Installation(ws, make_random)
     az_res_perm = AzureResourcePermissions(installation, ws, AzureResources(ws), location)
-    accounts = list(az_res_perm._get_storage_accounts())
-    assert len(accounts) == 1
-    assert accounts[0] == "labsazurethings"
+    az_res_perm.save_spn_permissions()
+
+    mapping = az_res_perm.load()
+    assert len(mapping) == 1
+    assert mapping[0].prefix == "labsazurethings"
 
 
 @pytest.mark.skip

tests/integration/workspace_access/test_groups.py

@@ -128,14 +128,11 @@ def test_delete_ws_groups_should_not_delete_non_reflected_acc_groups(ws, make_uc
 
 
 def validate_migrate_groups(group_manager: GroupManager, ws_group: Group, to_group: Group):
-    workspace_groups = group_manager._workspace_groups_in_workspace()
-    assert ws_group.display_name in workspace_groups
+    assert group_manager.has_workspace_group(ws_group.display_name), f'missing workspace group: {ws_group.display_name}'
     group_manager.rename_groups()
-    workspace_groups = group_manager._workspace_groups_in_workspace()
-    assert f"ucx-temp-{ws_group.display_name}" in workspace_groups
+    assert group_manager.has_workspace_group(f"ucx-temp-{ws_group.display_name}"), 'missing temp group'
     group_manager.reflect_account_groups_on_workspace()
-    account_workspace_groups = group_manager._account_groups_in_workspace()
-    assert to_group.display_name in account_workspace_groups
+    assert group_manager.has_account_group(to_group.display_name), f'missing account group: {to_group.display_name}'
 
 
 @retried(on=[NotFound], timeout=timedelta(minutes=2))

tests/integration/workspace_access/test_permissions_manager.py

@@ -10,7 +10,7 @@ def test_permissions_save_and_load(ws, sql_backend, inventory_schema, env_or_ski
         Permissions(object_id="efg", object_type="fgh", raw="ghi"),
     ]
 
-    pi._save(saved)
+    pi._save(saved)  # pylint: disable=protected-access
     loaded = pi.load_all()
 
     assert saved == loaded

tests/unit/assessment/test_azure.py

@@ -189,19 +189,9 @@ def test_list_all_pipeline_with_conf_spn_secret_avlb():
 
 
 def test_azure_spn_info_with_secret_unavailable():
-    ws = workspace_client_mock(cluster_ids=['simplest-autoscale'], pipeline_ids=['empty-spec'], secret_exists=False)
-    spark_conf = {
-        "spark.hadoop.fs.azure.account."
-        "oauth2.client.id.abcde.dfs.core.windows.net": "{{secrets/abcff/sp_app_client_id}}",
-        "spark.hadoop.fs.azure.account."
-        "oauth2.client.endpoint.abcde.dfs.core.windows.net": "https://login.microsoftonline.com/dedededede"
-        "/token",
-        "spark.hadoop.fs.azure.account."
-        "oauth2.client.secret.abcde.dfs.core.windows.net": "{{secrets/abcff/sp_secret}}",
-    }
-    crawler = AzureServicePrincipalCrawler(ws, MockBackend(), "ucx")._get_azure_spn_from_config(spark_conf)
-
-    assert crawler == set()
+    ws = workspace_client_mock(cluster_ids=['azure-spn-secret'], secret_exists=False)
+    result = AzureServicePrincipalCrawler(ws, MockBackend(), "ucx").snapshot()
+    assert len(result) == 0
 
 
 def test_jobs_assessment_with_spn_cluster_policy_not_found():

tests/unit/workspace_access/test_groups.py

@@ -539,24 +539,24 @@ def my_side_effect(group_id, **_):
         raise NotImplementedError
 
     wsclient.groups.get.side_effect = my_side_effect
+    wsclient.api_client.do.return_value = {
+        'Resources': [
+            {'displayName': 'group_1'},
+            {'displayName': 'group_2'},
+            {'displayName': 'group_3'},
+        ]
+    }
 
-    # Test when attributes do not contain "members"
     gm = GroupManager(backend, wsclient, inventory_database="inv")
-    result = gm._list_workspace_groups("WorkspaceGroup", "id,displayName,meta")
-    assert len(result) == 3
-    assert result[0].display_name == "group_1"
-    assert result[0].members is None
-    wsclient.groups.get.assert_not_called()
+    result = gm.snapshot()
 
-    # Test when attributes contain "members"
-    result = gm._list_workspace_groups("WorkspaceGroup", "id,displayName,meta,members")
     assert len(result) == 3
-    assert result[0].display_name == "group_1"
-    assert result[0].members == [
+    assert result[0].name_in_workspace == "group_1"
+    assert result[0].decode_members() == [
         ComplexValue(display="test-user-1", value="20"),
         ComplexValue(display="test-user-2", value="21"),
     ]
-    wsclient.groups.get.assert_called()
+    assert wsclient.groups.get.call_count == 3
 
 
 def test_snapshot_with_group_matched_by_suffix():

tests/unit/workspace_access/test_listing.py

@@ -52,7 +52,7 @@ def test_list_and_analyze_should_separate_folders_and_other_objects():
     client.workspace.list.return_value = [file, directory, notebook]
 
     listing_instance = listing.WorkspaceListing(client, 1)
-    directories, others = listing_instance._list_and_analyze(rootobj)
+    directories, others = listing_instance._list_and_analyze(rootobj)  # pylint: disable=protected-access
 
     assert compare(others, [file, notebook])
     assert compare(directories, [directory])

tests/unit/workspace_access/test_manager.py

@@ -33,7 +33,7 @@ def test_cleanup(b):
 def test_save(b):
     pi = PermissionManager(b, "test_database", [])
 
-    pi._save([Permissions("object1", "clusters", "test acl")])
+    pi._save([Permissions("object1", "clusters", "test acl")])  # pylint: disable=protected-access
 
     assert [Permissions(object_id="object1", object_type="clusters", raw="test acl")] == b.rows_written_for(
         "hive_metastore.test_database.permissions", "append"