|
| 1 | +import base64 |
1 | 2 | import json |
2 | 3 | import os |
3 | 4 | import tempfile |
@@ -257,6 +258,71 @@ def test_upload_with_integer_names(self): |
257 | 258 | f"{self.PROJECT_NAME}/test_folder", data_spec="multimodal" |
258 | 259 | ) |
259 | 260 |
|
| 261 | + def test_integration_not_found(self): |
| 262 | + with open(self.JSONL_ANNOTATIONS_PATH) as f: |
| 263 | + data = [json.loads(line) for line in f] |
| 264 | + with self.assertRaisesRegex(AppException, "Integration not found"): |
| 265 | + sa.upload_annotations( |
| 266 | + self.PROJECT_NAME, |
| 267 | + annotations=data, |
| 268 | + data_spec="multimodal", |
| 269 | + integration="non-existing-integration-xyz", |
| 270 | + ) |
| 271 | + |
| 272 | + def test_integration_only_supported_for_multimodal_data_spec(self): |
| 273 | + with open(self.JSONL_ANNOTATIONS_PATH) as f: |
| 274 | + data = [json.loads(line) for line in f] |
| 275 | + with self.assertRaisesRegex( |
| 276 | + AppException, "Integration is only supported for Multimodal projects" |
| 277 | + ): |
| 278 | + sa.upload_annotations( |
| 279 | + self.PROJECT_NAME, |
| 280 | + annotations=data, |
| 281 | + data_spec="default", |
| 282 | + integration="any-integration", |
| 283 | + ) |
| 284 | + |
| 285 | + def test_upload_with_existing_integration(self): |
| 286 | + integrations = sa.get_integrations() |
| 287 | + if not integrations: |
| 288 | + self.skipTest("No integrations available in the team.") |
| 289 | + integration = integrations[0] |
| 290 | + with open(self.JSONL_ANNOTATIONS_PATH) as f: |
| 291 | + data = [json.loads(line) for line in f] |
| 292 | + response = sa.upload_annotations( |
| 293 | + self.PROJECT_NAME, |
| 294 | + annotations=data, |
| 295 | + data_spec="multimodal", |
| 296 | + integration=integration["name"], |
| 297 | + ) |
| 298 | + assert len(response["succeeded"]) == 3 |
| 299 | + |
| 300 | + # Newly created items must carry the integration id used to sign URLs. |
| 301 | + # The item metadata's integration_id isn't exposed by the SDK entities, |
| 302 | + # so query the backend directly for it. |
| 303 | + from lib.core.jsx_conditions import EmptyQuery |
| 304 | + from lib.core.jsx_conditions import Join |
| 305 | + |
| 306 | + project, folder = sa.controller.get_project_folder( |
| 307 | + f"{self.PROJECT_NAME}/test_folder" |
| 308 | + ) |
| 309 | + item_service = sa.controller.service_provider.item_service |
| 310 | + client = item_service.client |
| 311 | + entity_context = base64.b64encode( |
| 312 | + f'{{"team_id":{client.team_id},"project_id":{project.id},' |
| 313 | + f'"folder_id":{folder.id}}}'.encode() |
| 314 | + ).decode() |
| 315 | + raw_items = client.jsx_paginate( |
| 316 | + url=item_service.URL_LIST, |
| 317 | + chunk_size=2000, |
| 318 | + body_query=EmptyQuery() & Join("metadata", ["path", "integration_id"]), |
| 319 | + method="post", |
| 320 | + headers={"x-sa-entity-context": entity_context}, |
| 321 | + ).data |
| 322 | + assert len(raw_items) == 3 |
| 323 | + for item in raw_items: |
| 324 | + assert item["metadata"]["integration_id"] == integration["id"] |
| 325 | + |
260 | 326 | def test_download_annotations(self): |
261 | 327 | with open(self.JSONL_ANNOTATIONS_PATH) as f: |
262 | 328 | data = [json.loads(line) for line in f] |
|
0 commit comments