Skip to content

Commit a11ae80

Browse files
HariGS-DBHariGS-DB
authored
added inventory table manager tests (#153)
1 parent 0d431c2 commit a11ae80

File tree

1 file changed

+208
-0
lines changed

1 file changed

+208
-0
lines changed

‎tests/unit/test_inventory_tables.py‎

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
1+
import json
2+
from unittest.mock import Mock
3+
4+
import pandas as pd
5+
import pytest
6+
from databricks.sdk.service.iam import AccessControlResponse, ObjectPermissions
7+
from pyspark.sql.types import StringType, StructField, StructType
8+
9+
from databricks.labs.ucx.config import InventoryConfig, InventoryTable
10+
from databricks.labs.ucx.inventory.table import InventoryTableManager
11+
from databricks.labs.ucx.inventory.types import (
12+
LogicalObjectType,
13+
PermissionsInventoryItem,
14+
RequestObjectType,
15+
)
16+
17+
18+
@pytest.fixture
19+
def workspace_client():
20+
client = Mock()
21+
return client
22+
23+
24+
perm_items = [PermissionsInventoryItem("object1", LogicalObjectType.CLUSTER, RequestObjectType.CLUSTERS, "test acl")]
25+
26+
27+
@pytest.fixture
28+
def inventory_table_manager(workspace_client, mocker):
29+
mocker.patch("databricks.labs.ucx.providers.spark.SparkMixin._initialize_spark", Mock())
30+
config = InventoryConfig(
31+
InventoryTable(catalog="test_catalog", database="test_database", name="test_inventory_table")
32+
)
33+
return InventoryTableManager(config, workspace_client)
34+
35+
36+
def test_inventory_table_manager_init(inventory_table_manager):
37+
assert str(inventory_table_manager.config.table) == "test_catalog.test_database.test_inventory_table"
38+
39+
40+
def test_table_schema(inventory_table_manager):
41+
schema = StructType(
42+
[
43+
StructField("object_id", StringType(), True),
44+
StructField("logical_object_type", StringType(), True),
45+
StructField("request_object_type", StringType(), True),
46+
StructField("raw_object_permissions", StringType(), True),
47+
]
48+
)
49+
assert inventory_table_manager._table_schema == schema
50+
51+
52+
def test_table(inventory_table_manager):
53+
assert inventory_table_manager._table == inventory_table_manager.spark.table(
54+
"test_catalog.test_database.test_inventory_table"
55+
)
56+
57+
58+
def test_cleanup(inventory_table_manager):
59+
inventory_table_manager.cleanup()
60+
inventory_table_manager.spark.sql.assert_called_with(
61+
"DROP TABLE IF EXISTS test_catalog.test_database.test_inventory_table"
62+
)
63+
64+
65+
def test_save(inventory_table_manager):
66+
inventory_table_manager.save(perm_items)
67+
inventory_table_manager.spark.createDataFrame.assert_called_once()
68+
69+
70+
def test_load_all(inventory_table_manager):
71+
items = pd.DataFrame(
72+
{
73+
"object_id": ["object1"],
74+
"logical_object_type": ["CLUSTER"],
75+
"request_object_type": ["clusters"],
76+
"raw_object_permissions": ["test acl"],
77+
}
78+
)
79+
inventory_table_manager._table.toPandas.return_value = items
80+
output = inventory_table_manager.load_all()
81+
assert output[0] == PermissionsInventoryItem(
82+
"object1", LogicalObjectType.CLUSTER, RequestObjectType.CLUSTERS, "test acl"
83+
)
84+
85+
86+
@pytest.mark.parametrize(
87+
"items,groups,status",
88+
[
89+
(
90+
PermissionsInventoryItem(
91+
object_id="group1",
92+
logical_object_type=LogicalObjectType.CLUSTER,
93+
request_object_type=RequestObjectType.CLUSTERS,
94+
raw_object_permissions=json.dumps(
95+
ObjectPermissions(
96+
object_id="clusterid1",
97+
object_type="clusters",
98+
access_control_list=[
99+
AccessControlResponse.from_dict(
100+
{"all_permissions": [{"permission_level": "CAN_MANAGE"}], "group_name": "group1"}
101+
),
102+
AccessControlResponse.from_dict(
103+
{"all_permissions": [{"permission_level": "CAN_MANAGE"}], "group_name": "admin"}
104+
),
105+
],
106+
).as_dict()
107+
),
108+
),
109+
["group1", "group2"],
110+
True,
111+
),
112+
(
113+
PermissionsInventoryItem(
114+
object_id="group1",
115+
logical_object_type=LogicalObjectType.ROLES,
116+
request_object_type=None,
117+
raw_object_permissions=json.dumps(
118+
{
119+
"roles": [
120+
{"value": "arn:aws:iam::123456789:instance-profile/test-uc-role"},
121+
{"value": "arn:aws:iam::123456789:instance-profile/test-uc-role2"},
122+
],
123+
"entitlements": [{"value": "workspace-access"}],
124+
}
125+
),
126+
),
127+
["group1", "group2"],
128+
True,
129+
),
130+
(
131+
PermissionsInventoryItem(
132+
object_id="scope-1",
133+
logical_object_type=LogicalObjectType.SECRET_SCOPE,
134+
request_object_type=None,
135+
raw_object_permissions="""{"acls": [
136+
{"principal": "g1", "permission": "READ"},
137+
{"principal": "unrelated-group", "permission": "READ"},
138+
{"principal": "admins", "permission": "MANAGE"}
139+
]}""",
140+
),
141+
["group1", "group2"],
142+
False,
143+
),
144+
(
145+
PermissionsInventoryItem(
146+
object_id="scope-1",
147+
logical_object_type=LogicalObjectType.SECRET_SCOPE,
148+
request_object_type=None,
149+
raw_object_permissions="""{"acls": [
150+
{"principal": "g1", "permission": "READ"},
151+
{"principal": "unrelated-group", "permission": "READ"},
152+
{"principal": "admins", "permission": "MANAGE"}
153+
]}""",
154+
),
155+
["g1", "group2"],
156+
True,
157+
),
158+
],
159+
)
160+
def test_is_item_relevant_to_groups(inventory_table_manager, items, groups, status):
161+
assert inventory_table_manager._is_item_relevant_to_groups(items, groups) is status
162+
163+
164+
def test_is_item_relevant_to_groups_exception(inventory_table_manager):
165+
item = PermissionsInventoryItem("object1", "FOO", "BAR", "test acl")
166+
with pytest.raises(NotImplementedError):
167+
inventory_table_manager._is_item_relevant_to_groups(item, ["g1"])
168+
169+
170+
def test_load_for_groups(inventory_table_manager):
171+
items = pd.DataFrame(
172+
{
173+
"object_id": ["group1"],
174+
"logical_object_type": ["CLUSTER"],
175+
"request_object_type": ["clusters"],
176+
"raw_object_permissions": json.dumps(
177+
ObjectPermissions(
178+
object_id="clusterid1",
179+
object_type="clusters",
180+
access_control_list=[
181+
AccessControlResponse.from_dict(
182+
{"all_permissions": [{"permission_level": "CAN_MANAGE"}], "group_name": "group1"}
183+
)
184+
],
185+
).as_dict()
186+
),
187+
}
188+
)
189+
groups = ["group1", "group2"]
190+
inventory_table_manager._table.toPandas.return_value = items
191+
output = inventory_table_manager.load_for_groups(groups)
192+
assert output[0] == PermissionsInventoryItem(
193+
object_id="group1",
194+
logical_object_type=LogicalObjectType.CLUSTER,
195+
request_object_type=RequestObjectType.CLUSTERS,
196+
raw_object_permissions=json.dumps(
197+
ObjectPermissions(
198+
object_id="clusterid1",
199+
object_type="clusters",
200+
access_control_list=[
201+
AccessControlResponse.from_dict(
202+
{"all_permissions": [{"permission_level": "CAN_MANAGE"}], "group_name": "group1"}
203+
)
204+
],
205+
).as_dict()
206+
),
207+
)
208+
assert len(output) == 1

0 commit comments

Comments
 (0)