Skip to content

Commit 4c3d935

Browse files
skyflow-srivyshnaviskyflow-srivyshnavi
committed
IE-509: refactored code
1 parent 8707925 commit 4c3d935

11 files changed

Lines changed: 386 additions & 101 deletions

‎tests/test_migrate_connections.py‎

Lines changed: 22 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,7 @@ def test_transform_connection_payload_removes_audit_and_invocation(monkeypatch):
1111
"name": "Conn",
1212
"mode": "EGRESS",
1313
"vaultID": "source-vault",
14-
"routes": [
15-
{"path": "/x", "method": "GET", "invocationURL": "https://example"}
16-
],
14+
"routes": [{"path": "/x", "method": "GET", "invocationURL": "https://example"}],
1715
"BasicAudit": {"CreatedBy": "u"},
1816
}
1917
out = mc.transform_connection_payload(source)
@@ -39,8 +37,8 @@ def test_main_with_config_creates_connection(mock_post, monkeypatch, tmp_path):
3937
sample = (
4038
"[\n"
4139
" {\n"
42-
" \"ID\": \"c1\", \"name\": \"Conn\", \"mode\": \"EGRESS\", \"vaultID\": \"sv\",\n"
43-
" \"routes\": [{\"path\": \"/p\", \"method\": \"GET\", \"invocationURL\": \"u\"}]\n"
40+
' "ID": "c1", "name": "Conn", "mode": "EGRESS", "vaultID": "sv",\n'
41+
' "routes": [{"path": "/p", "method": "GET", "invocationURL": "u"}]\n'
4442
" }\n"
4543
"]"
4644
)
@@ -150,21 +148,26 @@ def test_run_as_script_config_file(monkeypatch):
150148
monkeypatch.setenv("SOURCE_ACCOUNT_ID", "sacc")
151149
monkeypatch.setenv("SOURCE_ACCOUNT_AUTH", "stok")
152150
monkeypatch.setenv("SOURCE_ENV_URL", "https://s")
153-
sample = (
154-
"[\n {\n \"ID\": \"c1\", \"name\": \"Conn\", \"mode\": \"EGRESS\", \"vaultID\": \"sv\",\n \"routes\": [{\"path\": \"/p\", \"method\": \"GET\", \"invocationURL\": \"u\"}]\n }\n]"
155-
)
151+
sample = '[\n {\n "ID": "c1", "name": "Conn", "mode": "EGRESS", "vaultID": "sv",\n "routes": [{"path": "/p", "method": "GET", "invocationURL": "u"}]\n }\n]'
156152
with patch("builtins.open", mock_open(read_data=sample)):
157153
with patch.object(_requests, "post") as mpost:
158-
r = MagicMock(); r.status_code = 200; r.json.return_value = {"ID": "new"}
154+
r = MagicMock()
155+
r.status_code = 200
156+
r.json.return_value = {"ID": "new"}
159157
mpost.return_value = r
160158
runpy.run_module("migrate_connections", run_name="__main__")
161159
assert mpost.call_count >= 1
162160

161+
163162
@patch("migrate_connections.requests.get")
164163
def test_list_connections_combines_inbound_and_outbound(mock_get, monkeypatch):
165164
monkeypatch.setattr(mc, "SOURCE_ENV_URL", "https://s")
166-
r1 = MagicMock(); r1.raise_for_status.return_value = None; r1.json.return_value = {"ConnectionMappings": ["o1"]}
167-
r2 = MagicMock(); r2.raise_for_status.return_value = None; r2.json.return_value = {"ConnectionMappings": ["i1"]}
165+
r1 = MagicMock()
166+
r1.raise_for_status.return_value = None
167+
r1.json.return_value = {"ConnectionMappings": ["o1"]}
168+
r2 = MagicMock()
169+
r2.raise_for_status.return_value = None
170+
r2.json.return_value = {"ConnectionMappings": ["i1"]}
168171
mock_get.side_effect = [r1, r2]
169172
out = mc.list_connections("v")
170173
assert out == ["o1", "i1"]
@@ -173,7 +176,9 @@ def test_list_connections_combines_inbound_and_outbound(mock_get, monkeypatch):
173176
@patch("migrate_connections.requests.get")
174177
def test_get_connection_returns_json(mock_get, monkeypatch):
175178
monkeypatch.setattr(mc, "SOURCE_ENV_URL", "https://s")
176-
g = MagicMock(); g.raise_for_status.return_value = None; g.json.return_value = {"ID": "c"}
179+
g = MagicMock()
180+
g.raise_for_status.return_value = None
181+
g.json.return_value = {"ID": "c"}
177182
mock_get.return_value = g
178183
out = mc.get_connection("c")
179184
assert out["ID"] == "c"
@@ -186,14 +191,16 @@ def test_main_handles_creation_failure(mock_post, monkeypatch):
186191
monkeypatch.setattr(mc, "TARGET_VAULT_ID", "tv")
187192
monkeypatch.setattr(mc, "TARGET_ENV_URL", "https://t")
188193

189-
fail = MagicMock(); fail.status_code = 500; fail.content = b"err"
194+
fail = MagicMock()
195+
fail.status_code = 500
196+
fail.content = b"err"
190197
mock_post.return_value = fail
191198

192199
sample = (
193200
"[\n"
194201
" {\n"
195-
" \"ID\": \"c1\", \"name\": \"Conn\", \"mode\": \"INGRESS\", \"vaultID\": \"sv\",\n"
196-
" \"routes\": [{\"path\": \"/p\", \"method\": \"GET\", \"invocationURL\": \"u\"}]\n"
202+
' "ID": "c1", "name": "Conn", "mode": "INGRESS", "vaultID": "sv",\n'
203+
' "routes": [{"path": "/p", "method": "GET", "invocationURL": "u"}]\n'
197204
" }\n"
198205
"]"
199206
)

‎tests/test_migrate_policies.py‎

Lines changed: 36 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,9 @@ def test_transform_policy_payload_builds_rule_params(monkeypatch):
5555
assert len(out["ruleParams"]) == 3
5656
# Validate one conversion
5757
col_rule = out["ruleParams"][0]
58-
assert "columnRuleParams" in col_rule and col_rule["columnRuleParams"]["columns"][0].endswith("users.email")
58+
assert "columnRuleParams" in col_rule and col_rule["columnRuleParams"]["columns"][
59+
0
60+
].endswith("users.email")
5961

6062

6163
@patch("migrate_policies.requests.post")
@@ -114,7 +116,11 @@ def raise_err(_):
114116

115117
def test_main_generic_exception(monkeypatch):
116118
# Cause transform to fail with KeyError to hit generic except (111-113)
117-
monkeypatch.setattr(mp, "transform_policy_payload", lambda _: (_ for _ in ()).throw(Exception("boom")))
119+
monkeypatch.setattr(
120+
mp,
121+
"transform_policy_payload",
122+
lambda _: (_ for _ in ()).throw(Exception("boom")),
123+
)
118124
with pytest.raises(Exception):
119125
mp.main(policy_ids=["p1"])
120126

@@ -129,10 +135,35 @@ def test_run_as_script(monkeypatch):
129135
monkeypatch.setenv("SOURCE_ENV_URL", "https://s")
130136
monkeypatch.setenv("TARGET_ENV_URL", "https://t")
131137

132-
src = MagicMock(); src.raise_for_status.return_value = None; src.json.return_value = {"policy": {"ID": "p1", "namespace": "n", "status": "A", "BasicAudit": {}, "members": [], "rules": [{"ID": "r", "name": "R", "ruleExpression": "x", "actions": ["POLICY.read"], "resources": ["vault:v/table:t/column:c"], "resourceType": "COLUMN", "dlpFormat": None}]}}
133-
tgt = MagicMock(); tgt.raise_for_status.return_value = None; tgt.json.return_value = {"ID": "np"}
138+
src = MagicMock()
139+
src.raise_for_status.return_value = None
140+
src.json.return_value = {
141+
"policy": {
142+
"ID": "p1",
143+
"namespace": "n",
144+
"status": "A",
145+
"BasicAudit": {},
146+
"members": [],
147+
"rules": [
148+
{
149+
"ID": "r",
150+
"name": "R",
151+
"ruleExpression": "x",
152+
"actions": ["POLICY.read"],
153+
"resources": ["vault:v/table:t/column:c"],
154+
"resourceType": "COLUMN",
155+
"dlpFormat": None,
156+
}
157+
],
158+
}
159+
}
160+
tgt = MagicMock()
161+
tgt.raise_for_status.return_value = None
162+
tgt.json.return_value = {"ID": "np"}
134163

135-
with patch.object(_requests, "get", return_value=src) as mget, patch.object(_requests, "post", return_value=tgt) as mpost:
164+
with patch.object(_requests, "get", return_value=src) as mget, patch.object(
165+
_requests, "post", return_value=tgt
166+
) as mpost:
136167
runpy.run_module("migrate_policies", run_name="__main__")
137168
assert mget.call_count >= 1
138169
assert mpost.call_count >= 1

‎tests/test_migrate_roles.py‎

Lines changed: 69 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,8 @@ def test_transform_role_payload_filters_upstream(monkeypatch):
2222
out = mr.transform_role_payload(source)
2323
perms = out["roleDefinition"]["permissions"]
2424
assert "policies.read" in perms and all(
25-
p not in perms for p in [
25+
p not in perms
26+
for p in [
2627
"accounts.read:upstream",
2728
"workspaces.read:upstream",
2829
"vaults.read:upstream",
@@ -59,7 +60,9 @@ def test_main_system_role_path(mock_get, monkeypatch):
5960
@patch("migrate_roles.requests.post")
6061
@patch("migrate_roles.requests.get")
6162
@patch("migrate_roles.migrate_policies")
62-
def test_main_custom_role_create_and_assign(mock_migrate_policies, mock_get, mock_post, monkeypatch):
63+
def test_main_custom_role_create_and_assign(
64+
mock_migrate_policies, mock_get, mock_post, monkeypatch
65+
):
6366
monkeypatch.setattr(mr, "TARGET_VAULT_ID", "tv")
6467
monkeypatch.setattr(mr, "SOURCE_ENV_URL", "https://s")
6568
monkeypatch.setattr(mr, "TARGET_ENV_URL", "https://t")
@@ -74,9 +77,7 @@ def test_main_custom_role_create_and_assign(mock_migrate_policies, mock_get, moc
7477

7578
policies_resp = MagicMock()
7679
policies_resp.raise_for_status.return_value = None
77-
policies_resp.json.return_value = {
78-
"policies": [{"ID": "p1"}, {"ID": "p2"}]
79-
}
80+
policies_resp.json.return_value = {"policies": [{"ID": "p1"}, {"ID": "p2"}]}
8081

8182
# get_role -> role_resp, get_role_policies -> policies_resp
8283
mock_get.side_effect = [role_resp, policies_resp]
@@ -96,6 +97,7 @@ def test_main_custom_role_create_and_assign(mock_migrate_policies, mock_get, moc
9697
out = mr.main()
9798
assert out and any(r.get("ID") == "new-role" for r in out)
9899

100+
99101
@patch("migrate_roles.requests.post")
100102
@patch("migrate_roles.requests.get")
101103
def test_migrate_all_roles_branch(mock_get, mock_post, monkeypatch):
@@ -105,12 +107,22 @@ def test_migrate_all_roles_branch(mock_get, mock_post, monkeypatch):
105107
monkeypatch.setattr(mr, "SOURCE_ENV_URL", "https://s")
106108
monkeypatch.setattr(mr, "TARGET_ENV_URL", "https://t")
107109

108-
list_resp = MagicMock(); list_resp.raise_for_status.return_value = None; list_resp.json.return_value = {"roles": [{"ID": "r1"}]}
109-
role_resp = MagicMock(); role_resp.raise_for_status.return_value = None; role_resp.json.return_value = {"role": {"definition": {"name": "Custom", "permissions": ["policies.read"]}}}
110-
no_policies = MagicMock(); no_policies.raise_for_status.return_value = None; no_policies.json.return_value = {"policies": []}
110+
list_resp = MagicMock()
111+
list_resp.raise_for_status.return_value = None
112+
list_resp.json.return_value = {"roles": [{"ID": "r1"}]}
113+
role_resp = MagicMock()
114+
role_resp.raise_for_status.return_value = None
115+
role_resp.json.return_value = {
116+
"role": {"definition": {"name": "Custom", "permissions": ["policies.read"]}}
117+
}
118+
no_policies = MagicMock()
119+
no_policies.raise_for_status.return_value = None
120+
no_policies.json.return_value = {"policies": []}
111121
mock_get.side_effect = [list_resp, role_resp, no_policies]
112122

113-
create_post = MagicMock(); create_post.raise_for_status.return_value = None; create_post.json.return_value = {"ID": "new-role"}
123+
create_post = MagicMock()
124+
create_post.raise_for_status.return_value = None
125+
create_post.json.return_value = {"ID": "new-role"}
114126
mock_post.return_value = create_post
115127

116128
out = mr.main()
@@ -127,8 +139,14 @@ def test_skip_role_creation_if_exists(mock_get, mock_post, monkeypatch):
127139
monkeypatch.setattr(mr, "SOURCE_ENV_URL", "https://s")
128140
monkeypatch.setattr(mr, "TARGET_ENV_URL", "https://t")
129141

130-
role_resp = MagicMock(); role_resp.raise_for_status.return_value = None; role_resp.json.return_value = {"role": {"definition": {"name": "Custom", "permissions": ["policies.read"]}}}
131-
exists_resp = MagicMock(); exists_resp.raise_for_status.return_value = None; exists_resp.json.return_value = {"roles": [{"ID": "existing"}]}
142+
role_resp = MagicMock()
143+
role_resp.raise_for_status.return_value = None
144+
role_resp.json.return_value = {
145+
"role": {"definition": {"name": "Custom", "permissions": ["policies.read"]}}
146+
}
147+
exists_resp = MagicMock()
148+
exists_resp.raise_for_status.return_value = None
149+
exists_resp.json.return_value = {"roles": [{"ID": "existing"}]}
132150
mock_get.side_effect = [role_resp, exists_resp]
133151

134152
mr.main()
@@ -144,7 +162,10 @@ class Resp:
144162
err = requests.exceptions.HTTPError(response=Resp())
145163

146164
# First call returns a SYSTEM role so role_name is set
147-
role_resp = MagicMock(); role_resp.raise_for_status.return_value = None; role_resp.json.return_value = {"role": {"definition": {"name": mr.SYSTEM_ROLES[0]}}}
165+
role_resp = MagicMock()
166+
role_resp.raise_for_status.return_value = None
167+
role_resp.json.return_value = {"role": {"definition": {"name": mr.SYSTEM_ROLES[0]}}}
168+
148169
# get_system_role raises HTTPError
149170
def raise_err(*args, **kwargs):
150171
raise err
@@ -179,12 +200,22 @@ def test_custom_role_check_does_not_exist(mock_get, mock_post, monkeypatch):
179200
monkeypatch.setattr(mr, "SOURCE_ENV_URL", "https://s")
180201
monkeypatch.setattr(mr, "TARGET_ENV_URL", "https://t")
181202

182-
role_resp = MagicMock(); role_resp.raise_for_status.return_value = None; role_resp.json.return_value = {"role": {"definition": {"name": "Custom", "permissions": ["policies.read"]}}}
183-
no_exist = MagicMock(); no_exist.raise_for_status.return_value = None; no_exist.json.return_value = {"roles": []}
184-
policies_empty = MagicMock(); policies_empty.raise_for_status.return_value = None; policies_empty.json.return_value = {"policies": []}
203+
role_resp = MagicMock()
204+
role_resp.raise_for_status.return_value = None
205+
role_resp.json.return_value = {
206+
"role": {"definition": {"name": "Custom", "permissions": ["policies.read"]}}
207+
}
208+
no_exist = MagicMock()
209+
no_exist.raise_for_status.return_value = None
210+
no_exist.json.return_value = {"roles": []}
211+
policies_empty = MagicMock()
212+
policies_empty.raise_for_status.return_value = None
213+
policies_empty.json.return_value = {"policies": []}
185214
mock_get.side_effect = [role_resp, no_exist, policies_empty]
186215

187-
create_post = MagicMock(); create_post.raise_for_status.return_value = None; create_post.json.return_value = {"ID": "new-role"}
216+
create_post = MagicMock()
217+
create_post.raise_for_status.return_value = None
218+
create_post.json.return_value = {"ID": "new-role"}
188219
mock_post.return_value = create_post
189220

190221
out = mr.main()
@@ -197,19 +228,27 @@ def test_http_error_after_role_name(mock_get, mock_post, monkeypatch):
197228
# Raise HTTPError during create_role to hit lines 165-166
198229
class Resp:
199230
content = b"boom"
231+
200232
err = requests.exceptions.HTTPError(response=Resp())
201233

202234
monkeypatch.setattr(mr, "ROLE_IDS", "['rid']", raising=False)
203235
monkeypatch.setattr(mr, "TARGET_VAULT_ID", "tv")
204236
monkeypatch.setattr(mr, "SOURCE_ENV_URL", "https://s")
205237
monkeypatch.setattr(mr, "TARGET_ENV_URL", "https://t")
206238

207-
role_resp = MagicMock(); role_resp.raise_for_status.return_value = None; role_resp.json.return_value = {"role": {"definition": {"name": "Custom", "permissions": ["policies.read"]}}}
208-
policies_resp = MagicMock(); policies_resp.raise_for_status.return_value = None; policies_resp.json.return_value = {"policies": []}
239+
role_resp = MagicMock()
240+
role_resp.raise_for_status.return_value = None
241+
role_resp.json.return_value = {
242+
"role": {"definition": {"name": "Custom", "permissions": ["policies.read"]}}
243+
}
244+
policies_resp = MagicMock()
245+
policies_resp.raise_for_status.return_value = None
246+
policies_resp.json.return_value = {"policies": []}
209247
mock_get.side_effect = [role_resp, policies_resp]
210248

211249
def raise_err(*args, **kwargs):
212250
raise err
251+
213252
mock_post.side_effect = raise_err
214253

215254
with pytest.raises(requests.exceptions.HTTPError):
@@ -225,20 +264,29 @@ def test_generic_exception_after_role_name(mock_get, mock_post, monkeypatch):
225264
monkeypatch.setattr(mr, "SOURCE_ENV_URL", "https://s")
226265
monkeypatch.setattr(mr, "TARGET_ENV_URL", "https://t")
227266

228-
role_resp = MagicMock(); role_resp.raise_for_status.return_value = None; role_resp.json.return_value = {"role": {"definition": {"name": "Custom", "permissions": ["policies.read"]}}}
267+
role_resp = MagicMock()
268+
role_resp.raise_for_status.return_value = None
269+
role_resp.json.return_value = {
270+
"role": {"definition": {"name": "Custom", "permissions": ["policies.read"]}}
271+
}
229272
mock_get.side_effect = [role_resp]
230273

231-
create_post = MagicMock(); create_post.raise_for_status.return_value = None; create_post.json.return_value = {"ID": "new-role"}
274+
create_post = MagicMock()
275+
create_post.raise_for_status.return_value = None
276+
create_post.json.return_value = {"ID": "new-role"}
232277
mock_post.return_value = create_post
233278

234-
monkeypatch.setattr(mr, "get_role_policies", lambda _id: (_ for _ in ()).throw(Exception("oops")))
279+
monkeypatch.setattr(
280+
mr, "get_role_policies", lambda _id: (_ for _ in ()).throw(Exception("oops"))
281+
)
235282
with pytest.raises(Exception):
236283
mr.main()
237284

238285

239286
def test_run_as_script(monkeypatch):
240287
# Cover line 172
241288
import runpy, requests as _requests
289+
242290
monkeypatch.setenv("ROLE_IDS", "[]")
243291
monkeypatch.setenv("TARGET_VAULT_ID", "tv")
244292
monkeypatch.setenv("SOURCE_ENV_URL", "https://s")

0 commit comments

Comments
 (0)