Add folder assign test

Author: dshabin

superannotate/__init__.py

@@ -33,6 +33,7 @@ def consensus(*args, **kwargs):
     filter_annotation_instances, filter_images_by_comments,
     filter_images_by_tags
 )
+from .db.utils import _search_folders
 from .db.annotation_classes import (
     create_annotation_class, create_annotation_classes_from_classes_json,
     delete_annotation_class, download_annotation_classes_json,

superannotate/db/project_images.py

@@ -19,6 +19,7 @@
     get_project_default_image_quality_in_editor, _get_available_image_counts,
     get_project_metadata
 )
+from .teams import get_team_metadata
 from ..mixp.decorators import Trackable
 from .utils import _get_upload_auth_token, _get_boto_session_by_credentials, upload_image_array_to_s3, get_image_array_to_upload, __create_image, __copy_images, __move_images, get_project_folder_string
 
@@ -564,15 +565,14 @@ def assign_images(project, image_names, user):
     :param user: user email
     :type user: str
     """
-
-    # TODO:
-    # logger.info("Assign %s images to user %s", len(image_names), user)
-    # if len(image_names) == 0:
-    #     return
+    logger.info("Assign %s images to user %s", len(image_names), user)
+    if len(image_names) == 0:
+        return
 
     project, folder = get_project_and_folder_metadata(project)
-    if not folder:
-        folder = 'root'
+    folder_name = 'root'
+    if folder:
+        folder_name = folder['name']
 
     project_meta = get_project_metadata(project)
     params = {
@@ -582,7 +582,7 @@ def assign_images(project, image_names, user):
     json_req = {
         "image_names": image_names,
         "assign_user_id": user,
-        "folder_name": folder,
+        "folder_name": folder_name,
     }
     response = _api.send_request(
         req_type='PUT',
@@ -596,30 +596,6 @@ def assign_images(project, image_names, user):
             response.status_code, "Couldn't assign images " + response.text
         )
 
-    # images = search_images((project, project_folder), return_metadata=True)
-    # image_dict = {}
-    # for image in images:
-    #     image_dict[image["name"]] = image["id"]
-
-    # image_ids = []
-    # for image_name in image_names:
-    #     image_ids.append(image_dict[image_name])
-    # team_id, project_id = project["team_id"], project["id"]
-    # params = {"team_id": team_id, "project_id": project_id}
-    # if project_folder is not None:
-    #     params['folder_id'] = project_folder['id']
-    # json_req = {"user_id": user, "image_ids": image_ids}
-    # response = _api.send_request(
-    #     req_type='POST',
-    #     path='/images/assign',
-    #     params=params,
-    #     json_req=json_req
-    # )
-    # if not response.ok:
-    #     raise SABaseException(
-    #         response.status_code, "Couldn't assign images " + response.text
-    #     )
-
 
 def assign_folder(project, folder_name, users):
     """Assigns folder to users. With SDK, the user can be
@@ -634,9 +610,9 @@ def assign_folder(project, folder_name, users):
     """
 
     project_meta = get_project_metadata(project, include_contributors=True)
-    project_users = project_meta["contributors"]
+    project_users = get_team_metadata()["users"]
     project_name = project_meta['name']
-    project_users = [i['user_id'] for i in project_users]
+    project_users = [i['id'] for i in project_users]
     verified_contributor = []
 
     for user in users:
@@ -711,11 +687,20 @@ def unassign_images(project, image_names):
     :type image_names: list of str
     """
     project_meta = get_project_metadata(project)
+    project, folder = get_project_and_folder_metadata(project)
+    folder_name = 'root'
+    if folder:
+        folder_name = folder['name']
     params = {
         "project_id": project_meta['id'],
         "team_id": project_meta["team_id"]
     }
-    json_req = {"image_names": image_names, "remove_user_ids": ["all"]}
+    json_req = {
+        "image_names": image_names,
+        "remove_user_ids": ["all"],
+        "folder_name": folder_name
+    }
+
     response = _api.send_request(
         req_type='PUT',
         path='/images/editAssignment',

superannotate/db/utils.py

@@ -13,6 +13,7 @@
 from ..exceptions import SABaseException, SAImageSizeTooLarge, SANonExistingProjectNameException
 import datetime
 import boto3
+from .project_api import get_project_metadata_bare
 
 _api = API.get_instance()
 logger = logging.getLogger("superannotate-python-sdk")
@@ -731,3 +732,25 @@ def __attach_image_urls_to_project_thread(
         logger.warning(e)
     else:
         uploaded[thread_id] += uploaded_imgs
+
+
+def _search_folders(project, folder_name=None, includeUsers=False):
+    if not isinstance(project, dict):
+        project = get_project_metadata_bare(project)
+    team_id, project_id = project["team_id"], project["id"]
+    params = {
+        'team_id': team_id,
+        'project_id': project_id,
+        'offset': 0,
+        'name': folder_name,
+        'is_root': 0,
+        'includeUsers': includeUsers
+    }
+
+    response = _api.send_request(req_type='GET', path='/folders', params=params)
+    if not response.ok:
+        raise SABaseException(
+            response.status_code, "Couldn't search folders " + response.text
+        )
+    response = response.json()
+    return response

tests/test_assign_images.py

@@ -16,7 +16,7 @@ def test_assign_images(tmpdir):
     projects = sa.search_projects(PROJECT_NAME_VECTOR1, return_metadata=True)
     for project in projects:
         sa.delete_project(project)
-
+    time.sleep(1)
     project = sa.create_project(PROJECT_NAME_VECTOR1, "test", "Vector")
     email = sa.get_team_metadata()["users"][0]["email"]
     sa.share_project(project, email, "QA")
@@ -70,7 +70,7 @@ def test_assign_images_folder(tmpdir):
     projects = sa.search_projects(PROJECT_NAME_VECTOR2, return_metadata=True)
     for project in projects:
         sa.delete_project(project)
-
+    time.sleep(1)
     project = sa.create_project(PROJECT_NAME_VECTOR2, "test", "Vector")
     email = sa.get_team_metadata()["users"][0]["email"]
     sa.share_project(project, email, "QA")
@@ -118,4 +118,66 @@ def test_assign_images_folder(tmpdir):
     assert im1_metadata["annotator_id"] == email
     assert im2_metadata["annotator_id"] == email
     assert im1_metadata["qa_id"] is None
-    assert im2_metadata["qa_id"] is None
+    assert im2_metadata["qa_id"] is None
+
+
+def test_unassign_images(tmpdir):
+    tmpdir = Path(tmpdir)
+    projects = sa.search_projects(PROJECT_NAME_VECTOR1, return_metadata=True)
+    for project in projects:
+        sa.delete_project(project)
+    time.sleep(1)
+    project = sa.create_project(PROJECT_NAME_VECTOR1, "test", "Vector")
+    email = sa.get_team_metadata()["users"][0]["email"]
+    sa.share_project(project, email, "QA")
+    sa.upload_images_from_folder_to_project(
+        project, "./tests/sample_project_vector"
+    )
+    sa.assign_images(
+        project, ["example_image_1.jpg", "example_image_2.jpg"], email
+    )
+    sa.unassign_images(
+        project,
+        ["example_image_1.jpg", "example_image_2.jpg"],
+    )
+
+    im1_metadata = sa.get_image_metadata(project, "example_image_1.jpg")
+    im2_metadata = sa.get_image_metadata(project, "example_image_2.jpg")
+
+    assert im1_metadata["qa_id"] == None
+    assert im2_metadata["qa_id"] == None
+
+
+def test_assign_folder(tmpdir):
+    tmpdir = Path(tmpdir)
+    projects = sa.search_projects(PROJECT_NAME_VECTOR1, return_metadata=True)
+    for project in projects:
+        sa.delete_project(project)
+    time.sleep(1)
+    project = sa.create_project(PROJECT_NAME_VECTOR1, "test", "Vector")
+    folder_name = "assign_folder"
+    sa.create_folder(project, folder_name)
+    email = sa.get_team_metadata()["users"][1]["email"]
+    sa.share_project(project, email, "QA")
+    sa.assign_folder(project, folder_name, [email])
+    folders = sa._search_folders(PROJECT_NAME_VECTOR1, includeUsers=True)
+    assert len(folders["data"][0]['folder_users']) > 0
+
+
+def test_unassign_folder(tmpdir):
+    tmpdir = Path(tmpdir)
+    projects = sa.search_projects(PROJECT_NAME_VECTOR1, return_metadata=True)
+    for project in projects:
+        sa.delete_project(project)
+    time.sleep(1)
+    project = sa.create_project(PROJECT_NAME_VECTOR1, "test", "Vector")
+    folder_name = "assign_folder"
+    sa.create_folder(project, folder_name)
+    email = sa.get_team_metadata()["users"][1]["email"]
+    sa.share_project(project, email, "QA")
+    sa.assign_folder(project, folder_name, [email])
+    folders = sa._search_folders(PROJECT_NAME_VECTOR1, includeUsers=True)
+    assert len(folders["data"][0]['folder_users']) > 0
+    sa.unassign_folder(project, folder_name)
+    folders = sa._search_folders(PROJECT_NAME_VECTOR1, includeUsers=True)
+    assert len(folders["data"][0]['folder_users']) == 0