Skip to content

Commit baf3984

Browse files
nkvuongnkvuong
authored
Make azure unit tests readable again (#923)
## Changes Refactoring Azure unit tests to be more readable: - Expand `workspace_client_mock` to fetch job, pipeline, warehouse config from json - Expand `workspace_client_mock` to cover whether secret exists - Change `get_az_api_mapping` to fetch from json ### Tests <!-- How is this tested? Please see the checklist below and also describe any other relevant tests --> - [x] manually tested - [x] added unit tests - [ ] added integration tests - [ ] verified on staging environment (screenshot attached)
1 parent 2f1d8c4 commit baf3984

24 files changed

+793
-1564
lines changed

‎tests/unit/assessment/__init__.py‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,12 @@
33
from unittest.mock import create_autospec
44

55
from databricks.sdk import WorkspaceClient
6+
from databricks.sdk.errors import NotFound
67
from databricks.sdk.service.compute import ClusterDetails, Policy
8+
from databricks.sdk.service.jobs import BaseJob
9+
from databricks.sdk.service.pipelines import GetPipelineResponse, PipelineStateInfo
10+
from databricks.sdk.service.sql import EndpointConfPair
11+
from databricks.sdk.service.workspace import GetSecretResponse
712

813
__dir = pathlib.Path(__file__).parent
914

@@ -24,8 +29,43 @@ def _cluster_policy(policy_id: str):
2429
return Policy(description=definition, policy_family_definition_overrides=overrides)
2530

2631

27-
def workspace_client_mock(clusters="no-spark-conf.json"):
32+
def _pipeline(pipeline_id: str):
33+
fixture = _load_fixture(f"pipelines/{pipeline_id}.json")
34+
return GetPipelineResponse.from_dict(fixture)
35+
36+
37+
def _secret_not_found(secret_scope, secret_key):
38+
msg = f"Secret Scope {secret_scope} does not exist!"
39+
raise NotFound(msg)
40+
41+
42+
def get_az_api_mapping(*args, **kwargs):
43+
mapping = _load_fixture("../assessment/azure/mappings.json")[0]
44+
if args[1] in mapping:
45+
return mapping[args[1]]
46+
else:
47+
return {}
48+
49+
50+
def workspace_client_mock(
51+
clusters="no-spark-conf.json",
52+
pipelines="single-pipeline.json",
53+
pipeline_spec="empty-pipeline-spec.json",
54+
jobs="single-job.json",
55+
warehouse_config="single-config.json",
56+
secret_exists=True,
57+
):
2858
ws = create_autospec(WorkspaceClient)
2959
ws.clusters.list.return_value = _load_list(ClusterDetails, f"../assessment/clusters/{clusters}")
3060
ws.cluster_policies.get = _cluster_policy
61+
ws.pipelines.list_pipelines.return_value = _load_list(PipelineStateInfo, f"../assessment/pipelines/{pipelines}")
62+
ws.pipelines.get = _pipeline
63+
ws.jobs.list.return_value = _load_list(BaseJob, f"../assessment/jobs/{jobs}")
64+
ws.warehouses.get_workspace_warehouse_config().data_access_config = _load_list(
65+
EndpointConfPair, f"../assessment/warehouses/{warehouse_config}"
66+
)
67+
if secret_exists:
68+
ws.secrets.get_secret.return_value = GetSecretResponse(key="username", value="SGVsbG8sIFdvcmxkIQ==")
69+
else:
70+
ws.secrets.get_secret = _secret_not_found
3171
return ws

‎tests/unit/assessment/azure/mappings.json‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
[
2+
{
3+
"/v1.0/directoryObjects/user1": {
4+
"appId": "appIduser1",
5+
"displayName": "disNameuser1",
6+
"id": "Iduser1"
7+
},
8+
"/v1.0/directoryObjects/user2": {
9+
"appId": "appIduser2",
10+
"displayName": "disNameuser2",
11+
"id": "Iduser2"
12+
},
13+
"/v1.0/directoryObjects/user3": {
14+
"appId": "appIduser3",
15+
"displayName": "disNameuser3",
16+
"id": "Iduser3"
17+
},
18+
"/subscriptions": {
19+
"value": [
20+
{
21+
"displayName": "sub1",
22+
"subscriptionId": "001",
23+
"tenantId": "bar1"
24+
},
25+
{
26+
"displayName": "sub2",
27+
"subscriptionId": "002",
28+
"tenantId": "bar"
29+
},
30+
{
31+
"displayName": "sub3",
32+
"subscriptionId": "003",
33+
"tenantId": "bar3"
34+
}
35+
]
36+
},
37+
"/subscriptions/002/providers/Microsoft.Storage/storageAccounts": {
38+
"value": [
39+
{
40+
"name": "sto1"
41+
},
42+
{
43+
"name": "sto2",
44+
"id": "subscriptions/002/resourceGroups/rg1/storageAccounts/sto2"
45+
},
46+
{
47+
"name": "sto3",
48+
"id": "subscriptions/002/resourceGroups/rg1/storageAccounts/sto3"
49+
}
50+
]
51+
},
52+
"subscriptions/002/resourceGroups/rg1/storageAccounts/sto2/blobServices/default/containers": {
53+
"value": [
54+
{
55+
"name": "container1",
56+
"id": "subscriptions/002/resourceGroups/rg1/storageAccounts/sto2/containers/container1"
57+
},
58+
{
59+
"name": "container2",
60+
"id": "subscriptions/002/resourceGroups/rg1/storageAccounts/sto2/containers/container2"
61+
},
62+
{
63+
"name": "container3",
64+
"id": "subscriptions/002/resourceGroups/rg1/storageAccounts/sto2/containers/container3"
65+
},
66+
{
67+
"name": "container4"
68+
}
69+
]
70+
},
71+
"subscriptions/002/resourceGroups/rg1/storageAccounts/sto1/providers/Microsoft.Authorization/roleAssignments": {
72+
"value": [
73+
{
74+
"properties": {
75+
"principalId": "user1",
76+
"principalType": "User",
77+
"roleDefinitionId": "id001",
78+
"scope": "subscriptions/002/resourceGroups/rg1/storageAccounts/sto1"
79+
},
80+
"id": "rol1"
81+
}
82+
]
83+
},
84+
"subscriptions/002/resourceGroups/rg1/storageAccounts/sto2/providers/Microsoft.Authorization/roleAssignments": {
85+
"value": [
86+
{
87+
"properties": {
88+
"principalId": "user2",
89+
"principalType": "ServicePrincipal",
90+
"roleDefinitionId": "id001",
91+
"scope": "subscriptions/002/resourceGroups/rg1/storageAccounts/sto2"
92+
},
93+
"id": "rol1"
94+
}
95+
]
96+
},
97+
"id001": {
98+
"id": "role1",
99+
"properties": {
100+
"roleName": "Contributor"
101+
}
102+
},
103+
"subscriptions/002/resourceGroups/rg1/storageAccounts/sto3/providers/Microsoft.Authorization/roleAssignments": {
104+
"value": [
105+
{
106+
"properties": {
107+
"principalId": "user3",
108+
"principalType": "ServicePrincipal",
109+
"roleDefinitionId": "id002",
110+
"scope": "subscriptions/002/resourceGroups/rg1/storageAccounts/sto3"
111+
},
112+
"id": "rol1"
113+
},
114+
{
115+
"properties": {
116+
"principalId": "user3",
117+
"principalType": "ServicePrincipal",
118+
"roleDefinitionId": "id002",
119+
"scope": "subscriptions/002/resourceGroups/rg1/storageAccounts/sto3"
120+
},
121+
"id": "rol2"
122+
}
123+
]
124+
},
125+
"id002": {
126+
"id": "role2",
127+
"properties": {
128+
"roleName": "Storage Blob Data Owner"
129+
}
130+
},
131+
"subscriptions/002/resourceGroups/rg1/storageAccounts/sto2/containers/container1/providers/Microsoft.Authorization/roleAssignments": {
132+
"value": [
133+
{
134+
"properties": {
135+
"principalId": "user2",
136+
"principalType": "ServicePrincipal",
137+
"roleDefinitionId": "id001",
138+
"scope": "subscriptions/002/resourceGroups/rg1/storageAccounts/sto2/containers/container1"
139+
},
140+
"id": "rol1"
141+
}
142+
]
143+
},
144+
"subscriptions/002/resourceGroups/rg1/storageAccounts/sto2/containers/container2/providers/Microsoft.Authorization/roleAssignments": {
145+
"value": [
146+
{
147+
"properties": {
148+
"principalId": "user1",
149+
"principalType": "User",
150+
"roleDefinitionId": "id001",
151+
"scope": "this"
152+
},
153+
"id": "rol1"
154+
}
155+
]
156+
},
157+
"subscriptions/002/resourceGroups/rg1/storageAccounts/sto2/containers/container3/providers/Microsoft.Authorization/roleAssignments": {
158+
"value": [
159+
{
160+
"properties": {
161+
"principalId": "user3",
162+
"principalType": "ServicePrincipal",
163+
"roleDefinitionId": "id002",
164+
"scope": "subscriptions/002/resourceGroups/rg1/storageAccounts/sto2/containers/container3"
165+
},
166+
"id": "rol1"
167+
},
168+
{
169+
"properties": {
170+
"principalId": "user3",
171+
"principalType": "ServicePrincipal",
172+
"roleDefinitionId": "id002",
173+
"scope": "subscriptions/002/resourceGroups/rg1/storageAccounts/sto2"
174+
},
175+
"id": "rol2"
176+
}
177+
]
178+
}
179+
}
180+
]

‎tests/unit/assessment/clusters/assortment-spn.json‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
[
2+
{
3+
"autoscale": {
4+
"min_workers": 1,
5+
"max_workers": 6
6+
},
7+
"cluster_id": "0810-225833-atlanta69",
8+
"cluster_name": "Tech Summit FY24 Cluster-1",
9+
"cluster_source": "UI",
10+
"spark_context_id": 5134472582179565315,
11+
"spark_version": "9.3.x-cpu-ml-scala2.12"
12+
},
13+
{
14+
"autoscale": {
15+
"min_workers": 1,
16+
"max_workers": 6
17+
},
18+
"cluster_id": "0915-190044-3dqy6751",
19+
"cluster_name": "Tech Summit FY24 Cluster-2",
20+
"cluster_source": "UI",
21+
"spark_conf": {
22+
"spark.hadoop.fs.azure.account.oauth2.client.id.abcde.dfs.core.windows.net": "{{secrets/abcff/sp_app_client_id}}",
23+
"spark.hadoop.fs.azure.account.oauth2.client.endpoint.abcde.dfs.core.windows.net": "https://login.microsoftonline.com/dedededede/token",
24+
"spark.hadoop.fs.azure.account.oauth2.client.secret.abcde.dfs.core.windows.net": "{{secrets/abcff/sp_secret}}"
25+
},
26+
"spark_context_id": 5134472582179565315,
27+
"spark_version": "13.3.x-cpu-ml-scala2.12"
28+
}
29+
]

‎tests/unit/assessment/clusters/single-cluster-dupe-spn.json‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
[
2+
{
3+
"autoscale": {
4+
"max_workers": 6,
5+
"min_workers": 1
6+
},
7+
"spark_conf": {
8+
"spark.hadoop.fs.azure.account.oauth2.client.id.account.dfs.core.windows.net": "dummyclientidfromprofile",
9+
"spark.hadoop.fs.azure.account.oauth2.client.endpoint.account.dfs.core.windows.net": "https://login.microsoftonline.com/1234ededed/oauth2/token",
10+
"spark.hadoop.fs.azure.account.oauth2.client.secret.account.dfs.core.windows.net": "gfgfgfgfggfggfgfdds"
11+
},
12+
"cluster_source": "ClusterSource.UI",
13+
"spark_context_id": 5134472582179565315,
14+
"spark_version": "9.3.x-cpu-ml-scala2.12",
15+
"cluster_id": "0810-225833-atlanta69",
16+
"cluster_name": "Tech Summit FY24 Cluster-1"
17+
}
18+
]

‎tests/unit/assessment/clusters/single-cluster-spn-with-policy-overrides.json‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
[
2+
{
3+
"autoscale": {
4+
"max_workers": 6,
5+
"min_workers": 1
6+
},
7+
"cluster_source": "ClusterSource.UI",
8+
"spark_context_id": 5134472582179565315,
9+
"spark_version": "9.3.x-cpu-ml-scala2.12",
10+
"cluster_id": "0810-225833-atlanta69",
11+
"cluster_name": "Tech Summit FY24 Cluster-1",
12+
"policy_id": "spn-in-policy-overrides"
13+
}
14+
]

‎tests/unit/assessment/clusters/single-cluster-spn-with-policy.json‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
[
2+
{
3+
"autoscale": {
4+
"max_workers": 6,
5+
"min_workers": 1
6+
},
7+
"cluster_source": "ClusterSource.UI",
8+
"spark_context_id": 5134472582179565315,
9+
"spark_version": "9.3.x-cpu-ml-scala2.12",
10+
"cluster_id": "0810-225833-atlanta69",
11+
"cluster_name": "Tech Summit FY24 Cluster-1",
12+
"policy_id": "single-user-with-empty-appid-spn"
13+
}
14+
]

‎tests/unit/assessment/clusters/single-cluster-spn.json‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
[
2+
{
3+
"autoscale": {
4+
"max_workers": 6,
5+
"min_workers": 1
6+
},
7+
"spark_conf": {
8+
"spark.hadoop.fs.azure.account.oauth2.client.id.abcde.dfs.core.windows.net": "test123456789",
9+
"spark.hadoop.fs.azure.account.oauth2.client.endpoint.abcde.dfs.core.windows.net": "https://login.microsoftonline.com/dummy_tenant_id/oauth2/token",
10+
"spark.hadoop.fs.azure.account.oauth2.client.secret.abcde.dfs.core.windows.net": "{{secrets/abcff/sp_secret}}"
11+
},
12+
"cluster_source": "ClusterSource.UI",
13+
"spark_context_id": 5134472582179565315,
14+
"spark_version": "9.3.x-cpu-ml-scala2.12",
15+
"cluster_id": "0810-225833-atlanta69",
16+
"cluster_name": "Tech Summit FY24 Cluster-1"
17+
}
18+
]

0 commit comments

Comments
 (0)