Skip to content

Commit 9bb2be2

Browse files
committed
Test IAM Integration: remove after creation
1 parent 49d139c commit 9bb2be2

File tree

1 file changed

+128
-81
lines changed

1 file changed

+128
-81
lines changed

libs/labelbox/tests/integration/test_delegated_access.py

Lines changed: 128 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,18 @@
1313
)
1414

1515

16+
def delete_iam_integration(client, iam_integration_id: str):
17+
"""Helper function to delete an IAM integration using GraphQL mutation."""
18+
mutation = """mutation DeleteIamIntegrationPyApi($id: ID!) {
19+
deleteIamIntegration(where: { id: $id })
20+
}"""
21+
params = {
22+
"id": iam_integration_id
23+
}
24+
client.execute(mutation, params, experimental=True)
25+
26+
27+
1628
@pytest.fixture
1729
def test_integration_name() -> str:
1830
"""Returns a unique name for test integrations."""
@@ -33,7 +45,8 @@ def aws_integration(
3345
settings=settings,
3446
)
3547
yield integration
36-
# Cleanup will be handled by the organization's cleanup
48+
# Proper cleanup using delete mutation
49+
delete_iam_integration(client, integration.uid)
3750

3851

3952
@pytest.fixture
@@ -49,7 +62,8 @@ def gcp_integration(
4962
settings=settings,
5063
)
5164
yield integration
52-
# Cleanup will be handled by the organization's cleanup
65+
# Proper cleanup using delete mutation
66+
delete_iam_integration(client, integration.uid)
5367

5468

5569
@pytest.fixture
@@ -66,7 +80,8 @@ def azure_integration(
6680
settings=settings,
6781
)
6882
yield integration
69-
# Cleanup will be handled by the organization's cleanup
83+
# Proper cleanup using delete mutation
84+
delete_iam_integration(client, integration.uid)
7085

7186

7287
def test_create_aws_integration(client, test_integration_name):
@@ -79,11 +94,15 @@ def test_create_aws_integration(client, test_integration_name):
7994
name=test_integration_name, settings=settings
8095
)
8196

82-
assert integration.name == test_integration_name
83-
assert integration.provider == "AWS"
84-
assert isinstance(integration.settings, AwsIamIntegrationSettings)
85-
assert integration.settings.role_arn == settings.role_arn
86-
assert integration.settings.read_bucket == settings.read_bucket
97+
try:
98+
assert integration.name == test_integration_name
99+
assert integration.provider == "AWS"
100+
assert isinstance(integration.settings, AwsIamIntegrationSettings)
101+
assert integration.settings.role_arn == settings.role_arn
102+
assert integration.settings.read_bucket == settings.read_bucket
103+
finally:
104+
# Ensure cleanup even if assertions fail
105+
delete_iam_integration(client, integration.uid)
87106

88107

89108
def test_create_gcp_integration(client, test_integration_name):
@@ -93,10 +112,14 @@ def test_create_gcp_integration(client, test_integration_name):
93112
name=test_integration_name, settings=settings
94113
)
95114

96-
assert integration.name == test_integration_name
97-
assert integration.provider == "GCP"
98-
assert isinstance(integration.settings, GcpIamIntegrationSettings)
99-
assert integration.settings.read_bucket == settings.read_bucket
115+
try:
116+
assert integration.name == test_integration_name
117+
assert integration.provider == "GCP"
118+
assert isinstance(integration.settings, GcpIamIntegrationSettings)
119+
assert integration.settings.read_bucket == settings.read_bucket
120+
finally:
121+
# Ensure cleanup even if assertions fail
122+
delete_iam_integration(client, integration.uid)
100123

101124

102125
def test_create_azure_integration(client, test_integration_name):
@@ -109,13 +132,17 @@ def test_create_azure_integration(client, test_integration_name):
109132
name=test_integration_name, settings=settings
110133
)
111134

112-
assert integration.name == test_integration_name
113-
assert integration.provider == "Azure"
114-
assert isinstance(integration.settings, AzureIamIntegrationSettings)
115-
assert (
116-
integration.settings.read_container_url == settings.read_container_url
117-
)
118-
assert integration.settings.tenant_id == settings.tenant_id
135+
try:
136+
assert integration.name == test_integration_name
137+
assert integration.provider == "Azure"
138+
assert isinstance(integration.settings, AzureIamIntegrationSettings)
139+
assert (
140+
integration.settings.read_container_url == settings.read_container_url
141+
)
142+
assert integration.settings.tenant_id == settings.tenant_id
143+
finally:
144+
# Ensure cleanup even if assertions fail
145+
delete_iam_integration(client, integration.uid)
119146

120147

121148
def test_update_aws_integration(client, test_integration_name):
@@ -129,25 +156,29 @@ def test_update_aws_integration(client, test_integration_name):
129156
name=test_integration_name, settings=settings
130157
)
131158

132-
# Update integration
133-
new_settings = AwsIamIntegrationSettings(
134-
role_arn="arn:aws:iam::111111111111:role/updated",
135-
read_bucket="updated-bucket",
136-
)
137-
integration.update(
138-
name=f"updated-{test_integration_name}", settings=new_settings
139-
)
159+
try:
160+
# Update integration
161+
new_settings = AwsIamIntegrationSettings(
162+
role_arn="arn:aws:iam::111111111111:role/updated",
163+
read_bucket="updated-bucket",
164+
)
165+
integration.update(
166+
name=f"updated-{test_integration_name}", settings=new_settings
167+
)
140168

141-
# Verify update - find the specific integration by ID
142-
updated_integration = None
143-
for iam_int in client.get_organization().get_iam_integrations():
144-
if iam_int.uid == integration.uid:
145-
updated_integration = iam_int
146-
break
169+
# Verify update - find the specific integration by ID
170+
updated_integration = None
171+
for iam_int in client.get_organization().get_iam_integrations():
172+
if iam_int.uid == integration.uid:
173+
updated_integration = iam_int
174+
break
147175

148-
assert updated_integration is not None
149-
assert updated_integration.name == f"updated-{test_integration_name}"
150-
# Note: Settings may not be returned immediately after update
176+
assert updated_integration is not None
177+
assert updated_integration.name == f"updated-{test_integration_name}"
178+
# Note: Settings may not be returned immediately after update
179+
finally:
180+
# Ensure cleanup even if assertions fail
181+
delete_iam_integration(client, integration.uid)
151182

152183

153184
def test_update_gcp_integration(client, test_integration_name):
@@ -158,22 +189,26 @@ def test_update_gcp_integration(client, test_integration_name):
158189
name=test_integration_name, settings=settings
159190
)
160191

161-
# Update integration
162-
new_settings = GcpIamIntegrationSettings(read_bucket="gs://updated-bucket")
163-
integration.update(
164-
name=f"updated-{test_integration_name}", settings=new_settings
165-
)
192+
try:
193+
# Update integration
194+
new_settings = GcpIamIntegrationSettings(read_bucket="gs://updated-bucket")
195+
integration.update(
196+
name=f"updated-{test_integration_name}", settings=new_settings
197+
)
166198

167-
# Verify update - find the specific integration by ID
168-
updated_integration = None
169-
for iam_int in client.get_organization().get_iam_integrations():
170-
if iam_int.uid == integration.uid:
171-
updated_integration = iam_int
172-
break
199+
# Verify update - find the specific integration by ID
200+
updated_integration = None
201+
for iam_int in client.get_organization().get_iam_integrations():
202+
if iam_int.uid == integration.uid:
203+
updated_integration = iam_int
204+
break
173205

174-
assert updated_integration is not None
175-
assert updated_integration.name == f"updated-{test_integration_name}"
176-
# Note: Settings may not be returned immediately after update
206+
assert updated_integration is not None
207+
assert updated_integration.name == f"updated-{test_integration_name}"
208+
# Note: Settings may not be returned immediately after update
209+
finally:
210+
# Ensure cleanup even if assertions fail
211+
delete_iam_integration(client, integration.uid)
177212

178213

179214
def test_update_azure_integration(client, test_integration_name):
@@ -187,25 +222,29 @@ def test_update_azure_integration(client, test_integration_name):
187222
name=test_integration_name, settings=settings
188223
)
189224

190-
# Update integration
191-
new_settings = AzureIamIntegrationSettings(
192-
read_container_url="https://updated.blob.core.windows.net/test",
193-
tenant_id="updated-tenant",
194-
)
195-
integration.update(
196-
name=f"updated-{test_integration_name}", settings=new_settings
197-
)
225+
try:
226+
# Update integration
227+
new_settings = AzureIamIntegrationSettings(
228+
read_container_url="https://updated.blob.core.windows.net/test",
229+
tenant_id="updated-tenant",
230+
)
231+
integration.update(
232+
name=f"updated-{test_integration_name}", settings=new_settings
233+
)
198234

199-
# Verify update - find the specific integration by ID
200-
updated_integration = None
201-
for iam_int in client.get_organization().get_iam_integrations():
202-
if iam_int.uid == integration.uid:
203-
updated_integration = iam_int
204-
break
235+
# Verify update - find the specific integration by ID
236+
updated_integration = None
237+
for iam_int in client.get_organization().get_iam_integrations():
238+
if iam_int.uid == integration.uid:
239+
updated_integration = iam_int
240+
break
205241

206-
assert updated_integration is not None
207-
assert updated_integration.name == f"updated-{test_integration_name}"
208-
# Note: Settings may not be returned immediately after update
242+
assert updated_integration is not None
243+
assert updated_integration.name == f"updated-{test_integration_name}"
244+
# Note: Settings may not be returned immediately after update
245+
finally:
246+
# Ensure cleanup even if assertions fail
247+
delete_iam_integration(client, integration.uid)
209248

210249

211250
def test_update_azure_integration_with_credentials(
@@ -221,20 +260,24 @@ def test_update_azure_integration_with_credentials(
221260
name=test_integration_name, settings=settings
222261
)
223262

224-
# Update integration - note: credentials are not supported in updates
225-
new_settings = AzureIamIntegrationSettings(
226-
read_container_url="https://updated.blob.core.windows.net/test",
227-
tenant_id="updated-tenant",
228-
# Note: client_id and client_secret are not supported in update operations
229-
)
230-
integration.update(
231-
name=f"updated-{test_integration_name}", settings=new_settings
232-
)
263+
try:
264+
# Update integration - note: credentials are not supported in updates
265+
new_settings = AzureIamIntegrationSettings(
266+
read_container_url="https://updated.blob.core.windows.net/test",
267+
tenant_id="updated-tenant",
268+
# Note: client_id and client_secret are not supported in update operations
269+
)
270+
integration.update(
271+
name=f"updated-{test_integration_name}", settings=new_settings
272+
)
233273

234-
# Verify update (Note: credentials are not returned for security reasons)
235-
updated_integration = client.get_organization().get_iam_integrations()[0]
236-
assert updated_integration.name == f"updated-{test_integration_name}"
237-
# Note: Settings might not be returned for security reasons
274+
# Verify update (Note: credentials are not returned for security reasons)
275+
updated_integration = client.get_organization().get_iam_integrations()[0]
276+
assert updated_integration.name == f"updated-{test_integration_name}"
277+
# Note: Settings might not be returned for security reasons
278+
finally:
279+
# Ensure cleanup even if assertions fail
280+
delete_iam_integration(client, integration.uid)
238281

239282

240283
def test_validate_integration_format(client):
@@ -287,7 +330,8 @@ def test_set_as_default(client, test_integration_name):
287330
"""Test setting an integration as default."""
288331
# Save the original default integration
289332
original_default = client.get_organization().get_default_iam_integration()
290-
333+
334+
integration = None
291335
try:
292336
# Create an integration
293337
settings = AwsIamIntegrationSettings(
@@ -313,6 +357,9 @@ def test_set_as_default(client, test_integration_name):
313357
# Restore the original default integration
314358
if original_default is not None:
315359
original_default.set_as_default()
360+
# Clean up the created integration
361+
if integration is not None:
362+
delete_iam_integration(client, integration.uid)
316363

317364

318365
@pytest.mark.skip(

0 commit comments

Comments
 (0)