You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

97 lines
2.7 KiB

2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
  1. from typing import List
  2. from django.conf import settings
  3. from model_mommy import mommy
  4. from projects.models import Role, Member
  5. from projects.models import (
  6. DOCUMENT_CLASSIFICATION,
  7. SEQUENCE_LABELING,
  8. SEQ2SEQ,
  9. SPEECH2TEXT,
  10. IMAGE_CLASSIFICATION,
  11. INTENT_DETECTION_AND_SLOT_FILLING,
  12. )
  13. from roles.tests.utils import create_default_roles
  14. from users.tests.utils import make_user
  15. class ProjectData:
  16. def __init__(self, item, members):
  17. self.item = item
  18. self.members = members
  19. @property
  20. def admin(self):
  21. return self.members[0]
  22. @property
  23. def approver(self):
  24. return self.members[1]
  25. @property
  26. def annotator(self):
  27. return self.members[2]
  28. @property
  29. def staffs(self):
  30. return [self.approver, self.annotator]
  31. def assign_user_to_role(project_member, project, role_name):
  32. role, _ = Role.objects.get_or_create(name=role_name)
  33. if Member.objects.filter(user=project_member, project=project).exists():
  34. mapping = Member.objects.get(user=project_member, project=project)
  35. mapping.role = role
  36. mapping.save()
  37. else:
  38. mapping = Member.objects.get_or_create(role_id=role.id, user_id=project_member.id, project_id=project.id)
  39. return mapping
  40. def make_project(task: str, users: List[str], roles: List[str], collaborative_annotation=False, **kwargs):
  41. create_default_roles()
  42. # create users.
  43. users = [make_user(name) for name in users]
  44. # create a project.
  45. project_model = {
  46. DOCUMENT_CLASSIFICATION: "TextClassificationProject",
  47. SEQUENCE_LABELING: "SequenceLabelingProject",
  48. SEQ2SEQ: "Seq2seqProject",
  49. SPEECH2TEXT: "Speech2TextProject",
  50. IMAGE_CLASSIFICATION: "ImageClassificationProject",
  51. INTENT_DETECTION_AND_SLOT_FILLING: "IntentDetectionAndSlotFillingProject",
  52. }.get(task, "Project")
  53. project = mommy.make(
  54. _model=project_model,
  55. project_type=task,
  56. collaborative_annotation=collaborative_annotation,
  57. created_by=users[0],
  58. **kwargs,
  59. )
  60. # assign roles to the users.
  61. for user, role in zip(users, roles):
  62. assign_user_to_role(user, project, role)
  63. return ProjectData(item=project, members=users)
  64. def make_tag(project):
  65. return mommy.make("Tag", project=project)
  66. def prepare_project(task: str = "Any", collaborative_annotation=False, **kwargs):
  67. return make_project(
  68. task=task,
  69. users=["admin", "approver", "annotator"],
  70. roles=[
  71. settings.ROLE_PROJECT_ADMIN,
  72. settings.ROLE_ANNOTATION_APPROVER,
  73. settings.ROLE_ANNOTATOR,
  74. ],
  75. collaborative_annotation=collaborative_annotation,
  76. **kwargs,
  77. )