Skip to content

Commit 16753c7

Browse files
paulnoirelpaulnoirel
committed
Split method and rename it. Split tests
1 parent f6f2d54 commit 16753c7

File tree

2 files changed

+164
-78
lines changed

2 files changed

+164
-78
lines changed

‎libs/labelbox/src/labelbox/schema/dataset.py‎

Lines changed: 94 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -914,79 +914,103 @@ def _convert_items_to_upsert_format(self, _items):
914914
_upsert_items.append(DataRowUpsertItem(payload=item, id=key))
915915
return _upsert_items
916916

917-
def set_iam_integration(self, iam_integration: Union[str, IAMIntegration, None] = None) -> Optional[IAMIntegration]:
918-
"""
917+
def add_iam_integration(self, iam_integration: Union[str, IAMIntegration]) -> IAMIntegration:
918+
"""
919919
Sets the IAM integration for the dataset. IAM integration is used to sign URLs for data row assets.
920920
921-
>>> # Get all IAM integrations
922-
>>> iam_integrations = client.get_organization().get_iam_integrations()
923-
>>>
924-
>>> # Get IAM integration id
925-
>>> iam_integration_id = [integration.uid for integration
926-
>>> in iam_integrations
927-
>>> if integration.name == "My S3 integration"][0]
928-
>>>
929-
>>> # Set IAM integration for integration id
930-
>>> dataset.set_iam_integration(iam_integration_id)
931-
>>>
932-
>>> # Get IAM integration object
933-
>>> iam_integration = [integration.uid for integration
934-
>>> in iam_integrations
935-
>>> if integration.name == "My S3 integration"][0]
936-
>>>
937-
>>> # Set IAM integration for IAMIntegrtion object
938-
>>> dataset.set_iam_integration(iam_integration)
939-
>>>
940-
>>> # Unset IAM integration
941-
>>> dataset.set_iam_integration()
921+
Args:
922+
iam_integration (Union[str, IAMIntegration]): IAM integration object or IAM integration id.
923+
924+
Returns:
925+
IAMIntegration: IAM integration object.
926+
927+
Raises:
928+
LabelboxError: If the IAM integration can't be set.
929+
930+
Examples:
931+
932+
>>> # Get all IAM integrations
933+
>>> iam_integrations = client.get_organization().get_iam_integrations()
934+
>>>
935+
>>> # Get IAM integration id
936+
>>> iam_integration_id = [integration.uid for integration
937+
>>> in iam_integrations
938+
>>> if integration.name == "My S3 integration"][0]
939+
>>>
940+
>>> # Set IAM integration for integration id
941+
>>> dataset.set_iam_integration(iam_integration_id)
942+
>>>
943+
>>> # Get IAM integration object
944+
>>> iam_integration = [integration.uid for integration
945+
>>> in iam_integrations
946+
>>> if integration.name == "My S3 integration"][0]
947+
>>>
948+
>>> # Set IAM integration for IAMIntegrtion object
949+
>>> dataset.set_iam_integration(iam_integration)
942950
"""
943951

944-
# Unset IAM integration if iam_integration is None
945-
if iam_integration is None:
946-
query = """mutation DetachSignerPyApi($id: ID!) {
947-
clearSignerForDataset(where: {id: $id}) {
948-
id
949-
signer {
950-
id
951-
}
952-
}
953-
}"""
954-
response = self.client.execute(query, {"id": self.uid})
955-
956-
if response:
957-
return response["clearSignerForDataset"]["signer"]
958-
else:
959-
raise LabelboxError("Can't unset IAM integration")
952+
iam_integration_id = iam_integration.uid if isinstance(iam_integration, IAMIntegration) else iam_integration
953+
954+
query = """
955+
mutation SetSignerForDatasetPyApi($signerId: ID!, $datasetId: ID!) {
956+
setSignerForDataset(
957+
data: { signerId: $signerId }
958+
where: { id: $datasetId }
959+
) {
960+
id
961+
signer {
962+
id
963+
}
964+
}
965+
}
966+
"""
967+
968+
response = self.client.execute(query, {"signerId": iam_integration_id, "datasetId": self.uid})
969+
970+
if not response:
971+
raise LabelboxError(f"Can't set IAM integration {iam_integration_id}")
972+
973+
try:
974+
iam_integration_id = response.get("setSignerForDataset", {}).get("signer", {})["id"]
975+
return [integration for integration
976+
in self.client.get_organization().get_iam_integrations()
977+
if integration.uid == iam_integration_id][0]
978+
except:
979+
raise LabelboxError(f"Can't retrieve IAM integration {iam_integration_id}")
960980

961-
else:
962-
963-
if isinstance(iam_integration, IAMIntegration):
964-
iam_integration_id = iam_integration.uid
965-
else:
966-
iam_integration_id = iam_integration
967-
968-
query = """mutation AttachSignerPyApi($signerId: ID!, $datasetId: ID!) {
969-
setSignerForDataset(data: {signerId: $signerId}, where: {id: $datasetId}) {
970-
id
971-
signer {
972-
id
973-
}
974-
}
975-
}"""
976-
response = self.client.execute(query, {"signerId": iam_integration_id, "datasetId": self.uid})
981+
def remove_iam_integration(self) -> None:
982+
"""
983+
Unsets the IAM integration for the dataset.
984+
985+
Args:
986+
None
987+
988+
Returns:
989+
None
990+
991+
Raises:
992+
LabelboxError: If the IAM integration can't be unset.
993+
994+
Examples:
995+
>>> dataset.remove_iam_integration()
996+
"""
997+
998+
query = """
999+
mutation DetachSignerPyApi($id: ID!) {
1000+
clearSignerForDataset(where: { id: $id }) {
1001+
id
1002+
signer {
1003+
id
1004+
}
1005+
}
1006+
}
1007+
"""
1008+
1009+
response = self.client.execute(query, {"id": self.uid})
1010+
1011+
if response:
1012+
return response.get("clearSignerForDataset", {}).get("signer")
1013+
else:
1014+
raise LabelboxError("Can't unset IAM integration")
9771015

978-
# Return IAM Integration object if
979-
if response:
980-
try:
981-
982-
iam_integration_id = response.get("setSignerForDataset", {}).get("signer", {})["id"]
983-
return [integration for integration
984-
in self.client.get_organization().get_iam_integrations()
985-
if integration.uid == iam_integration_id][0]
986-
except:
987-
raise LabelboxError(f"Can't retrieve IAM integration {iam_integration_id}")
988-
989-
else:
990-
raise LabelboxError(f"Can't set IAM integration {iam_integration_id}")
991-
992-
return response
1016+

‎libs/labelbox/tests/integration/test_delegated_access.py‎

Lines changed: 70 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ def test_no_default_integration(client):
8686
)
8787
@pytest.mark.skipif(not os.environ.get("DA_GCP_LABELBOX_API_KEY"),
8888
reason="DA_GCP_LABELBOX_API_KEY not found")
89-
def test_integration_change():
89+
def test_add_integration_from_object():
9090
"""
9191
This test is based on test_non_default_integration() and assumes the following:
9292
@@ -106,24 +106,86 @@ def test_integration_change():
106106
in integrations
107107
if 'aws-da-test-bucket' in integration.name][0]
108108

109-
ds = client.create_dataset(iam_integration=integration, name=f"integration_change-{uuid.uuid4()}")
110-
111-
# Test unset integration
112-
ds.set_iam_integration()
113-
assert ds.iam_integration() is None
109+
ds = client.create_dataset(iam_integration=None, name=f"integration_change-{uuid.uuid4()}")
114110

115111
# Test set integration with object
116-
new_integration = ds.set_iam_integration(integration)
112+
new_integration = ds.add_iam_integration(integration)
117113
assert new_integration == integration
118114

115+
# Cleaning
116+
ds.delete()
117+
118+
@pytest.mark.skip(
119+
reason=
120+
"Google credentials are being updated for this test, disabling till it's all sorted out"
121+
)
122+
@pytest.mark.skipif(not os.environ.get("DA_GCP_LABELBOX_API_KEY"),
123+
reason="DA_GCP_LABELBOX_API_KEY not found")
124+
def test_add_integration_from_uid():
125+
"""
126+
This test is based on test_non_default_integration() and assumes the following:
127+
128+
1. aws delegated access is configured to work with lbox-test-bucket
129+
2. an integration called aws is available to the org
130+
131+
Currently tests against:
132+
Org ID: cl26d06tk0gch10901m7jeg9v
133+
Email: jtso+aws_sdk_tests@labelbox.com
134+
"""
135+
client = Client(api_key=os.environ.get("DA_GCP_LABELBOX_API_KEY"))
136+
integrations = client.get_organization().get_iam_integrations()
137+
138+
# Prepare dataset with an existing integration
139+
integration = [
140+
integration for integration
141+
in integrations
142+
if 'aws-da-test-bucket' in integration.name][0]
143+
144+
ds = client.create_dataset(iam_integration=None, name=f"integration_change-{uuid.uuid4()}")
145+
119146
# Test set integration with integration id
120147
integration_id = [
121148
integration.uid for integration
122149
in integrations
123150
if 'aws-da-test-bucket' in integration.name][0]
124151

125-
new_integration = ds.set_iam_integration(integration_id)
152+
new_integration = ds.add_iam_integration(integration_id)
126153
assert new_integration == integration
127154

155+
# Cleaning
156+
ds.delete()
157+
158+
@pytest.mark.skip(
159+
reason=
160+
"Google credentials are being updated for this test, disabling till it's all sorted out"
161+
)
162+
@pytest.mark.skipif(not os.environ.get("DA_GCP_LABELBOX_API_KEY"),
163+
reason="DA_GCP_LABELBOX_API_KEY not found")
164+
def test_integration_remove():
165+
"""
166+
This test is based on test_non_default_integration() and assumes the following:
167+
168+
1. aws delegated access is configured to work with lbox-test-bucket
169+
2. an integration called aws is available to the org
170+
171+
Currently tests against:
172+
Org ID: cl26d06tk0gch10901m7jeg9v
173+
Email: jtso+aws_sdk_tests@labelbox.com
174+
"""
175+
client = Client(api_key=os.environ.get("DA_GCP_LABELBOX_API_KEY"))
176+
integrations = client.get_organization().get_iam_integrations()
177+
178+
# Prepare dataset with an existing integration
179+
integration = [
180+
integration for integration
181+
in integrations
182+
if 'aws-da-test-bucket' in integration.name][0]
183+
184+
ds = client.create_dataset(iam_integration=integration, name=f"integration_change-{uuid.uuid4()}")
185+
186+
# Test unset integration
187+
ds.remove_iam_integration()
188+
assert ds.iam_integration() is None
189+
128190
# Cleaning
129191
ds.delete()

0 commit comments

Comments
 (0)