|
|
from typing import List
from django.conf import settings from model_mommy import mommy
from projects.models import Role, Member from projects.models import ( DOCUMENT_CLASSIFICATION, SEQUENCE_LABELING, SEQ2SEQ, SPEECH2TEXT, IMAGE_CLASSIFICATION, INTENT_DETECTION_AND_SLOT_FILLING, ) from roles.tests.utils import create_default_roles from users.tests.utils import make_user
class ProjectData: def __init__(self, item, members): self.item = item self.members = members
@property def admin(self): return self.members[0]
@property def approver(self): return self.members[1]
@property def annotator(self): return self.members[2]
@property def staffs(self): return [self.approver, self.annotator]
def assign_user_to_role(project_member, project, role_name): role, _ = Role.objects.get_or_create(name=role_name) if Member.objects.filter(user=project_member, project=project).exists(): mapping = Member.objects.get(user=project_member, project=project) mapping.role = role mapping.save() else: mapping = Member.objects.get_or_create(role_id=role.id, user_id=project_member.id, project_id=project.id) return mapping
def make_project(task: str, users: List[str], roles: List[str], collaborative_annotation=False, **kwargs): create_default_roles()
# create users. users = [make_user(name) for name in users]
# create a project. project_model = { DOCUMENT_CLASSIFICATION: "TextClassificationProject", SEQUENCE_LABELING: "SequenceLabelingProject", SEQ2SEQ: "Seq2seqProject", SPEECH2TEXT: "Speech2TextProject", IMAGE_CLASSIFICATION: "ImageClassificationProject", INTENT_DETECTION_AND_SLOT_FILLING: "IntentDetectionAndSlotFillingProject", }.get(task, "Project") project = mommy.make( _model=project_model, project_type=task, collaborative_annotation=collaborative_annotation, created_by=users[0], **kwargs, )
# assign roles to the users. for user, role in zip(users, roles): assign_user_to_role(user, project, role)
return ProjectData(item=project, members=users)
def make_tag(project): return mommy.make("Tag", project=project)
def prepare_project(task: str = "Any", collaborative_annotation=False, **kwargs): return make_project( task=task, users=["admin", "approver", "annotator"], roles=[ settings.ROLE_PROJECT_ADMIN, settings.ROLE_ANNOTATION_APPROVER, settings.ROLE_ANNOTATOR, ], collaborative_annotation=collaborative_annotation, **kwargs, )
|