From ca7f114e158691065cb7f39a383b723d3b9699b3 Mon Sep 17 00:00:00 2001 From: Val Brodsky Date: Thu, 26 Sep 2024 12:18:29 -0700 Subject: [PATCH] Add wait for labels Simplify test for user group creation / deletion Rewrite annotation upload tests for LLM projects --- libs/labelbox/tests/conftest.py | 66 +++++++++++++------ .../test_generic_data_types.py | 52 +++++++++++++-- .../integration/schema/test_user_group.py | 32 +++------ libs/labelbox/tests/integration/test_label.py | 4 +- 4 files changed, 103 insertions(+), 51 deletions(-) diff --git a/libs/labelbox/tests/conftest.py b/libs/labelbox/tests/conftest.py index 6d13a8d83..49eab165d 100644 --- a/libs/labelbox/tests/conftest.py +++ b/libs/labelbox/tests/conftest.py @@ -1,35 +1,39 @@ -from datetime import datetime -from random import randint -from string import ascii_letters - import json import os import re -import uuid import time -from labelbox.schema.project import Project -import requests -from labelbox.schema.ontology import Ontology -import pytest -from types import SimpleNamespace -from typing import Type +import uuid +from datetime import datetime from enum import Enum -from typing import Tuple +from random import randint +from string import ascii_letters +from types import SimpleNamespace +from typing import Tuple, Type + +import pytest +import requests -from labelbox import Dataset, DataRow -from labelbox import MediaType +from labelbox import ( + Classification, + Client, + DataRow, + Dataset, + LabelingFrontend, + MediaType, + OntologyBuilder, + Option, + Tool, +) +from labelbox.exceptions import LabelboxError from labelbox.orm import query from labelbox.pagination import PaginatedCollection +from labelbox.schema.annotation_import import LabelImport +from labelbox.schema.enums import AnnotationImportState from labelbox.schema.invite import Invite +from labelbox.schema.ontology import Ontology +from labelbox.schema.project import Project from labelbox.schema.quality_mode import QualityMode from labelbox.schema.queue_mode import QueueMode -from labelbox import Client - -from labelbox import LabelingFrontend -from labelbox import OntologyBuilder, Tool, Option, Classification -from labelbox.schema.annotation_import import LabelImport -from labelbox.schema.enums import AnnotationImportState -from labelbox.exceptions import LabelboxError IMG_URL = "https://picsum.photos/200/300.jpg" MASKABLE_IMG_URL = "https://storage.googleapis.com/labelbox-datasets/image_sample_data/2560px-Kitano_Street_Kobe01s5s4110.jpeg" @@ -1255,6 +1259,21 @@ def teardown_ontology_feature_schemas(ontology: Ontology): class ModuleTearDownHelpers(TearDownHelpers): ... +class LabelHelpers: + def wait_for_labels(self, project, number_of_labels=1): + timeout_seconds = 10 + while True: + labels = list(project.labels()) + if len(labels) >= number_of_labels: + return labels + timeout_seconds -= 2 + if timeout_seconds <= 0: + raise TimeoutError( + f"Timed out waiting for label for project '{project.uid}' to finish processing" + ) + time.sleep(2) + + @pytest.fixture def teardown_helpers(): return TearDownHelpers() @@ -1263,3 +1282,8 @@ def teardown_helpers(): @pytest.fixture(scope="module") def module_teardown_helpers(): return TearDownHelpers() + + +@pytest.fixture +def label_helpers(): + return LabelHelpers() diff --git a/libs/labelbox/tests/data/annotation_import/test_generic_data_types.py b/libs/labelbox/tests/data/annotation_import/test_generic_data_types.py index 76236f0dd..1cc5538d9 100644 --- a/libs/labelbox/tests/data/annotation_import/test_generic_data_types.py +++ b/libs/labelbox/tests/data/annotation_import/test_generic_data_types.py @@ -112,11 +112,6 @@ def test_generic_data_row_type_by_global_key( (MediaType.Conversational, MediaType.Conversational), (MediaType.Document, MediaType.Document), (MediaType.Dicom, MediaType.Dicom), - ( - MediaType.LLMPromptResponseCreation, - MediaType.LLMPromptResponseCreation, - ), - (MediaType.LLMPromptCreation, MediaType.LLMPromptCreation), (OntologyKind.ResponseCreation, OntologyKind.ResponseCreation), (OntologyKind.ModelEvaluation, OntologyKind.ModelEvaluation), ], @@ -186,6 +181,53 @@ def test_import_media_types( assert exported_annotations == expected_data +@pytest.mark.parametrize( + "configured_project, media_type", + [ + ( + MediaType.LLMPromptResponseCreation, + MediaType.LLMPromptResponseCreation, + ), + (MediaType.LLMPromptCreation, MediaType.LLMPromptCreation), + ], + indirect=["configured_project"], +) +def test_import_media_types_llm( + client: Client, + configured_project: Project, + annotations_by_media_type, + exports_v2_by_media_type, + export_v2_test_helpers, + helpers, + media_type, + wait_for_label_processing, +): + annotations_ndjson = list( + itertools.chain.from_iterable(annotations_by_media_type[media_type]) + ) + + label_import = lb.LabelImport.create_from_objects( + client, + configured_project.uid, + f"test-import-{media_type}", + annotations_ndjson, + ) + label_import.wait_until_done() + + assert label_import.state == AnnotationImportState.FINISHED + assert len(label_import.errors) == 0 + + all_annotations = sorted([a["uuid"] for a in annotations_ndjson]) + successful_annotations = sorted( + [ + status["uuid"] + for status in label_import.statuses + if status["status"] == "SUCCESS" + ] + ) + assert successful_annotations == all_annotations + + @pytest.mark.parametrize( "configured_project_by_global_key, media_type", [ diff --git a/libs/labelbox/tests/integration/schema/test_user_group.py b/libs/labelbox/tests/integration/schema/test_user_group.py index 6aebd4e89..50aaad4a7 100644 --- a/libs/labelbox/tests/integration/schema/test_user_group.py +++ b/libs/labelbox/tests/integration/schema/test_user_group.py @@ -1,13 +1,13 @@ -import pytest -import faker from uuid import uuid4 -from labelbox import Client -from labelbox.schema.user_group import UserGroup, UserGroupColor + +import faker +import pytest + from labelbox.exceptions import ( - ResourceNotFoundError, ResourceCreationError, - UnprocessableEntityError, + ResourceNotFoundError, ) +from labelbox.schema.user_group import UserGroup, UserGroupColor data = faker.Faker() @@ -147,9 +147,6 @@ def test_cannot_update_group_id(user_group): def test_get_user_groups_with_creation_deletion(client): user_group = None try: - # Get all user groups - user_groups = list(UserGroup(client).get_user_groups()) - # manual delete for iterators group_name = data.name() user_group = UserGroup(client) @@ -157,25 +154,12 @@ def test_get_user_groups_with_creation_deletion(client): user_group.create() user_groups_post_creation = list(UserGroup(client).get_user_groups()) + assert user_group in user_groups_post_creation - # Verify that at least one user group is returned - assert len(user_groups_post_creation) > 0 - assert len(user_groups_post_creation) == len(user_groups) + 1 - - # Verify that each user group has a valid ID and name - for ug in user_groups_post_creation: - assert ug.id is not None - assert ug.name is not None - - user_group.delete() user_group = None user_groups_post_deletion = list(UserGroup(client).get_user_groups()) - - assert ( - len(user_groups_post_deletion) == len(user_groups_post_creation) - 1 - ) - + assert user_group not in user_groups_post_deletion finally: if user_group: user_group.delete() diff --git a/libs/labelbox/tests/integration/test_label.py b/libs/labelbox/tests/integration/test_label.py index 1bd8a8276..0daa4758e 100644 --- a/libs/labelbox/tests/integration/test_label.py +++ b/libs/labelbox/tests/integration/test_label.py @@ -41,11 +41,13 @@ def test_label_update(configured_project_with_label): assert label.label == "something else" -def test_label_filter_order(configured_project_with_label): +def test_label_filter_order(configured_project_with_label, label_helpers): project, _, _, label = configured_project_with_label l1 = label project.create_label() + label_helpers.wait_for_labels(project, 2) + l2 = next(project.labels()) assert set(project.labels()) == {l1, l2}