BDD tests for removing project in training and testing dataset

interactive_ai/tests/e2e/features/project_removal.feature

@@ -26,3 +26,14 @@ Feature: project removal
       Then the request is rejected
       When the user tries to load the image 'cat.jpg'
       Then the request is rejected
+
+  Scenario: deletion of a project in training
+    Given an annotated project of type 'detection'
+    And the user requests to train a model
+    And a job of type 'train' is running
+    When the user tries to delete the project
+    Then the request is rejected
+    When the user cancels the job
+    And the user deletes the project
+    And the user tries to load the project
+    Then the request is rejected

interactive_ai/tests/e2e/features/steps/annotation_management.py

@@ -81,7 +81,10 @@ def _annotate_image_or_frame(  # noqa: C901, PLR0912, PLR0915
     media_name: str,
     labels: list[str],
     frame_index: int | None = None,
+    dataset_id: str | None = None,
 ) -> None:
+    if dataset_id is None:
+        dataset_id = context.dataset_id
     annotations_api: AnnotationsApi = context.annotations_api
 
     media_details = context._media_info_by_name[media_name]
@@ -277,7 +280,7 @@ def _annotate_image_or_frame(  # noqa: C901, PLR0912, PLR0915
                 organization_id=context.organization_id,
                 workspace_id=context.workspace_id,
                 project_id=context.project_id,
-                dataset_id=context.dataset_id,
+                dataset_id=dataset_id,
                 image_id=media_details.id,
                 create_image_annotation_request=CreateImageAnnotationRequest(
                     media_identifier=media_identifier,
@@ -289,7 +292,7 @@ def _annotate_image_or_frame(  # noqa: C901, PLR0912, PLR0915
                 organization_id=context.organization_id,
                 workspace_id=context.workspace_id,
                 project_id=context.project_id,
-                dataset_id=context.dataset_id,
+                dataset_id=dataset_id,
                 video_id=media_details.id,
                 frame_index=frame_index,
                 create_video_frame_annotation_request=CreateVideoFrameAnnotationRequest(
@@ -302,17 +305,19 @@ def _annotate_image_or_frame(  # noqa: C901, PLR0912, PLR0915
 
 
 def _get_image_or_frame_annotations(
-    context: Context, media_type: str, media_name: str, frame_index: int | None = None
+    context: Context, media_type: str, media_name: str, frame_index: int | None = None, dataset_id: str | None = None
 ) -> list[Annotation]:
     annotations_api: AnnotationsApi = context.annotations_api
     annotations: list[Annotation] = []
+    if dataset_id is None:
+        dataset_id = context.dataset_id
 
     if media_type == "image":
         ann_response = annotations_api.get_image_annotation(
             organization_id=context.organization_id,
             workspace_id=context.workspace_id,
             project_id=context.project_id,
-            dataset_id=context.dataset_id,
+            dataset_id=dataset_id,
             image_id=context._media_info_by_name[media_name].id,
             annotation_id="latest",
         )
@@ -321,7 +326,7 @@ def _get_image_or_frame_annotations(
             organization_id=context.organization_id,
             workspace_id=context.workspace_id,
             project_id=context.project_id,
-            dataset_id=context.dataset_id,
+            dataset_id=dataset_id,
             video_id=context._media_info_by_name[media_name].id,
             frame_index=frame_index,
             annotation_id="latest",
@@ -335,6 +340,19 @@ def _get_image_or_frame_annotations(
     return annotations
 
 
+@when("the user annotates the image '{image_name}' in dataset '{dataset_name}' with label '{label_name}'")
+def step_when_user_annotates_image_in_dataset_with_custom_label(
+    context: Context, image_name: str, dataset_name: str, label_name: str
+) -> None:
+    _annotate_image_or_frame(
+        context=context,
+        media_type="image",
+        media_name=image_name,
+        labels=[label_name],
+        dataset_id=context._dataset_info_by_name[dataset_name].id,
+    )
+
+
 @when("the user annotates the image '{image_name}' with label '{label_name}'")
 def step_when_user_annotates_image_with_custom_label(context: Context, image_name: str, label_name: str) -> None:
     _annotate_image_or_frame(context=context, media_type="image", media_name=image_name, labels=[label_name])
@@ -372,6 +390,43 @@ def step_when_user_tries_annotating_image_with_custom_labels(
         context.exception = e
 
 
+@then("the image '{image_name}' in dataset '{dataset_name}' has labels '{raw_expected_label_names}'")
+def step_then_image_in_dataset_has_expected_labels(
+    context: Context, image_name: str, dataset_name: str, raw_expected_label_names: str
+) -> None:
+    annotations = _get_image_or_frame_annotations(
+        context=context,
+        media_type="image",
+        media_name=image_name,
+        dataset_id=context._dataset_info_by_name[dataset_name].id,
+    )
+
+    expected_label_names = raw_expected_label_names.split(", ")
+
+    label_info_by_id = {label.id: label for label in context.label_info_by_name.values()}
+    found_label_names = [
+        label_info_by_id[label_id].name for annotation in annotations for label_id in annotation.label_ids
+    ]
+    expected_label_names_freq = Counter(expected_label_names)
+    found_label_names_freq = Counter(found_label_names)
+    assert expected_label_names_freq == found_label_names_freq, (
+        f"Expected to find labels with the respective frequency: {expected_label_names_freq}, "
+        f"found instead: {found_label_names_freq}"
+    )
+
+
+@then("the image '{image_name}' in dataset '{dataset_name}' has label '{expected_label_name}'")
+def step_then_image_in_dataset_has_expected_label(
+    context: Context, image_name: str, dataset_name: str, expected_label_name: str
+) -> None:
+    step_then_image_in_dataset_has_expected_labels(
+        context=context,
+        image_name=image_name,
+        dataset_name=dataset_name,
+        raw_expected_label_names=expected_label_name,
+    )
+
+
 @then("the image '{image_name}' has labels '{raw_expected_label_names}'")
 def step_then_image_has_expected_labels(context: Context, image_name: str, raw_expected_label_names: str) -> None:
     annotations = _get_image_or_frame_annotations(context=context, media_type="image", media_name=image_name)

interactive_ai/tests/e2e/features/steps/job_monitoring.py

@@ -16,7 +16,7 @@
 import time
 from typing import TYPE_CHECKING
 
-from behave import then
+from behave import step
 from behave.runner import Context
 from static_definitions import JobState, JobType
 
@@ -29,7 +29,7 @@
 logger = logging.getLogger(__name__)
 
 
-@then("a job of type '{job_type:w}' is scheduled")
+@step("a job of type '{job_type:w}' is scheduled")
 def step_then_job_scheduled(context: Context, job_type: str) -> None:
     """Asserts that the "job_type" is scheduled"""
     jobs_api: JobsApi = context.jobs_api
@@ -45,7 +45,42 @@ def step_then_job_scheduled(context: Context, job_type: str) -> None:
     assert found_job_type == expected_job_type, f"Expected job type {expected_job_type}, but found {found_job_type}"
 
 
-@then("the job completes successfully within {job_timeout:d} minutes")
+@step("a job of type '{job_type:w}' is running")
+def step_then_job_running(context: Context, job_type: str) -> None:
+    """Asserts that the "job_type" is running"""
+    jobs_api: JobsApi = context.jobs_api
+    expected_job_type = JobType(job_type)
+
+    for _ in range(20):
+        context.job_info = jobs_api.get_job(
+            organization_id=context.organization_id,
+            workspace_id=context.workspace_id,
+            job_id=context.job_id,
+        )
+        if context.job_info.actual_instance.state == JobState.RUNNING:
+            time.sleep(10)  # provide some time for the job to properly start
+            break
+        time.sleep(1)
+    else:
+        raise RuntimeError(f"A job of type '{job_type}' is not running within 20 seconds")
+
+    found_job_type = JobType(context.job_info.actual_instance.type)
+    assert found_job_type == expected_job_type, f"Expected job type {expected_job_type}, but found {found_job_type}"
+
+
+@step("the user cancels the job")
+def step_when_user_cancels_job(context: Context) -> None:
+    """Cancels the job"""
+    jobs_api: JobsApi = context.jobs_api
+    jobs_api.cancel_job(
+        organization_id=context.organization_id,
+        workspace_id=context.workspace_id,
+        job_id=context.job_id,
+    )
+    time.sleep(5)  # provide some time for the job to stop
+
+
+@step("the job completes successfully within {job_timeout:d} minutes")
 def step_then_job_finishes_before_timeout(context: Context, job_timeout: int) -> None:
     """
     Asserts that the jobs finishes within the passed "job_timeout". If not, raises a TimeoutError.

interactive_ai/tests/e2e/features/steps/media_management.py

@@ -18,7 +18,7 @@
 
 import cv2
 import numpy as np
-from behave import given, when
+from behave import given, then, when
 from behave.runner import Context
 
 if TYPE_CHECKING:
@@ -28,7 +28,7 @@
 logger = logging.getLogger(__name__)
 
 
-def _create_and_upload_random_image(context: Context, image_name: str) -> None:
+def _create_and_upload_random_image(context: Context, image_name: str, dataset_id: str) -> None:
     random_image = np.random.randint(0, 256, (100, 100, 3), dtype=np.uint8)
     img = Image.fromarray(random_image)
 
@@ -42,15 +42,15 @@ def _create_and_upload_random_image(context: Context, image_name: str) -> None:
             organization_id=context.organization_id,
             workspace_id=context.workspace_id,
             project_id=context.project_id,
-            dataset_id=context.dataset_id,
+            dataset_id=dataset_id,
             file=image_path,  # file.read(),
             _headers={"Content-Disposition": f'form-data; name="file"; filename="{image_name}"'},
         )
 
         context._media_info_by_name[image_name] = upload_image_response
 
 
-def _create_and_upload_random_video(context: Context, video_name: str) -> None:
+def _create_and_upload_random_video(context: Context, video_name: str, dataset_id: str) -> None:
     with tempfile.TemporaryDirectory() as temp_dir:
         video_path = f"{temp_dir}/{video_name}"
 
@@ -71,7 +71,7 @@ def _create_and_upload_random_video(context: Context, video_name: str) -> None:
             organization_id=context.organization_id,
             workspace_id=context.workspace_id,
             project_id=context.project_id,
-            dataset_id=context.dataset_id,
+            dataset_id=dataset_id,
             file=video_path,
             _headers={"Content-Disposition": f'form-data; name="file"; filename="{video_name}"'},
         )
@@ -81,12 +81,12 @@ def _create_and_upload_random_video(context: Context, video_name: str) -> None:
 
 @given("an image called '{image_name}'")
 def step_given_image_with_custom_name(context: Context, image_name: str) -> None:
-    _create_and_upload_random_image(context=context, image_name=image_name)
+    _create_and_upload_random_image(context=context, image_name=image_name, dataset_id=context.dataset_id)
 
 
 @given("a video called '{video_name}'")
 def step_given_video_with_custom_name(context: Context, video_name: str) -> None:
-    _create_and_upload_random_video(context=context, video_name=video_name)
+    _create_and_upload_random_video(context=context, video_name=video_name, dataset_id=context.dataset_id)
 
 
 @when("the user loads the image '{image_name}'")
@@ -101,9 +101,27 @@ def step_when_user_loads_image(context: Context, image_name: str) -> None:
     )
 
 
+@when("the user uploads an image called '{image_name}' to dataset '{dataset_name}'")
+def step_when_user_uploads_image_to_dataset(context: Context, image_name: str, dataset_name: str) -> None:
+    dataset_id = context._dataset_info_by_name[dataset_name].id
+    _create_and_upload_random_image(context=context, image_name=image_name, dataset_id=dataset_id)
+
+
 @when("the user tries to load the image '{image_name}'")
 def step_when_user_tries_loading_image(context: Context, image_name: str) -> None:
     try:
         step_when_user_loads_image(context=context, image_name=image_name)
     except Exception as e:
         context.exception = e
+
+
+@then("the dataset '{dataset_name}' contains an image called '{image_name}'")
+def step_then_dataset_contains_image(context: Context, dataset_name: str, image_name: str) -> None:
+    media_api: MediaApi = context.media_api
+    context._media_info_by_name[image_name] = media_api.get_image_detail(
+        organization_id=context.organization_id,
+        workspace_id=context.workspace_id,
+        project_id=context.project_id,
+        dataset_id=context._dataset_info_by_name[dataset_name].id,
+        image_id=context._media_info_by_name[image_name].id,
+    )

interactive_ai/tests/e2e/features/steps/project_management.py

@@ -271,6 +271,9 @@ def _create_project_from_scratch(  # noqa: C901
     context.project_info = create_project_response  # store the whole response for convenience in other steps
     context.project_id = create_project_response.id
     context.dataset_id = create_project_response.datasets[0].id
+    context._dataset_info_by_name = {
+        dataset_info.name: dataset_info for dataset_info in create_project_response.datasets
+    }
     context.label_info_by_name = {  # dict[str, CreateProject201ResponsePipelineTasksInnerLabelsInner]
         label.name: label for task in create_project_response.pipeline.tasks if task.labels for label in task.labels
     }
@@ -364,6 +367,14 @@ def step_when_user_loads_project(context: Context) -> None:
     )
 
 
+@when("the user tries to delete the project")
+def step_when_user_tries_to_delete_project(context: Context) -> None:
+    try:
+        step_when_user_deletes_project(context=context)
+    except Exception as e:
+        context.exception = e
+
+
 @when("the user tries to load the project")
 def step_when_user_tries_to_load_project(context: Context) -> None:
     try:

interactive_ai/tests/e2e/features/steps/training.py

@@ -12,7 +12,7 @@
 # with no express or implied warranties, other than those that are expressly stated
 # in the License.
 
-from behave import then, when
+from behave import step, then, when
 from behave.runner import Context
 from geti_client import ModelsApi, TrainJob, TrainModelRequest
 from static_definitions import TaskType
@@ -49,7 +49,7 @@ def _train(
     context.job_id = train_response.job_id
 
 
-@when("the user requests to train a model")
+@step("the user requests to train a model")
 def step_when_user_trains_model(context: Context):
     _train(context=context)