From 29f00b35f4296cc5d9c726446bfdca8485631cd6 Mon Sep 17 00:00:00 2001
From: Cor <jczuurmond@protonmail.com>
Date: Wed, 18 Sep 2024 15:02:18 +0200
Subject: [PATCH] Let `validate-groups-membership` command to run as collection
 (#2657)

## Changes
Let `validate-groups-membership` command to run as collection

### Linked issues

Resolves #2613

### Functionality

- [x] modified existing command: `databricks labs ucx
validate-groups-membership`

### Tests

- [x] manually tested
- [x] added unit tests
- [ ] ~added integration tests~ : Covering after #2507
---
 labs.yml                       |  3 +++
 src/databricks/labs/ucx/cli.py | 17 +++++++++++++----
 tests/unit/test_cli.py         | 18 +++++++++++++++---
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/labs.yml b/labs.yml
index 117c6e69f2..13e272b0f6 100644
--- a/labs.yml
+++ b/labs.yml
@@ -178,6 +178,9 @@ commands:
       Workspace Group Name\tMembers Count\tAccount Group Name\tMembers Count\tDifference
       {{range .}}{{.wf_group_name}}\t{{.wf_group_members_count}}\t{{.acc_group_name}}\t{{.acc_group_members_count}}\t{{.group_members_difference}}
       {{end}}
+    flags:
+    - name: run-as-collection
+      description: (Optional) Run the command for the collection of workspaces with ucx installed. Default is False.
 
   - name: migrate-credentials
     description: Migrate credentials for storage access to UC storage credential
diff --git a/src/databricks/labs/ucx/cli.py b/src/databricks/labs/ucx/cli.py
index d1b9ede53a..1b3e106f53 100644
--- a/src/databricks/labs/ucx/cli.py
+++ b/src/databricks/labs/ucx/cli.py
@@ -203,11 +203,20 @@ def repair_run(w: WorkspaceClient, step):
 
 
 @ucx.command
-def validate_groups_membership(w: WorkspaceClient):
+def validate_groups_membership(
+    w: WorkspaceClient,
+    ctx: WorkspaceContext | None = None,
+    run_as_collection: bool = False,
+    a: AccountClient | None = None,
+) -> None:
     """Validate the groups to see if the groups at account level and workspace level has different membership"""
-    ctx = WorkspaceContext(w)
-    mismatch_groups = ctx.group_manager.validate_group_membership()
-    print(json.dumps(mismatch_groups))
+    if ctx:
+        workspace_contexts = [ctx]
+    else:
+        workspace_contexts = _get_workspace_contexts(w, a, run_as_collection)
+    for workspace_context in workspace_contexts:
+        mismatch_groups = workspace_context.group_manager.validate_group_membership()
+        print(json.dumps(mismatch_groups))
 
 
 @ucx.command
diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py
index 74b99818b2..b969053d97 100644
--- a/tests/unit/test_cli.py
+++ b/tests/unit/test_cli.py
@@ -407,9 +407,21 @@ def test_save_storage_and_principal_azure(ws, caplog, acc_client):
     azure_resource_permissions.save_spn_permissions.assert_called_once()
 
 
-def test_validate_groups_membership(ws):
-    validate_groups_membership(ws)
-    ws.groups.list.assert_called()
+@pytest.mark.parametrize("run_as_collection", [True, False])
+def test_validate_groups_membership_lists_groups(
+    run_as_collection,
+    workspace_clients,
+    acc_client,
+) -> None:
+    if not run_as_collection:
+        workspace_clients = [workspace_clients[0]]
+    validate_groups_membership(
+        workspace_clients[0],
+        run_as_collection=run_as_collection,
+        a=acc_client,
+    )
+    for workspace_client in workspace_clients:
+        workspace_client.groups.list.assert_called()
 
 
 def test_save_storage_and_principal_aws(ws, acc_client):