system-pri/backend/tests/fake_data.py

from typing import List

from .factory import ProjectSupervisorFactory, YearGroupProjectSupervisorsFactory, \
    StudentFactory, YearGroupStudentsFactory, GroupFactory
from app.dependencies import db
from app.project_supervisor.models import YearGroup, ProjectSupervisor
from app.students.models import Group, Student, YearGroupStudents
from app.base.mode import ModeGroups


def create_year_group(data: dict = None) -> YearGroup:
    if data is None:
        data = {'mode': ModeGroups.STATIONARY.value, 'name': '2022/2023'}
    yg = YearGroup(**data)
    db.session.add(yg)
    db.session.commit()
    return yg


def create_project_supervisors(yg: YearGroup, amount: int, limit_group: int = 3) -> List[ProjectSupervisorFactory]:
    ps = [ProjectSupervisorFactory() for _ in range(amount)]
    db.session.add_all(ps)
    db.session.commit()
    db.session.add_all(
        [YearGroupProjectSupervisorsFactory(limit_group=limit_group, year_group_id=yg.id, project_supervisor_id=p.id)
         for p in ps])
    db.session.commit()
    return ps


def create_dummy_ps(data: dict) -> ProjectSupervisor:
    ps = ProjectSupervisor(**data)
    db.session.add(ps)
    db.session.commit()
    return ps


def create_dummy_group(data: dict, project_supervisor_id: int) -> Group:
    group = Group(**data, project_supervisor_id=project_supervisor_id)
    db.session.add(group)
    db.session.commit()
    return group


def create_students(yg: YearGroup, amount: int) -> List[StudentFactory]:
    students = [StudentFactory() for _ in range(amount)]
    db.session.add_all(students)
    db.session.commit()
    db.session.add_all([YearGroupStudentsFactory(year_group_id=yg.id, student_index=s.index) for s in students])
    db.session.commit()
    return students


def create_student(data: dict, year_group_id: int = None) -> Student:
    st = Student(**data)
    db.session.add(st)
    db.session.commit()
    if year_group_id is not None:
        db.session.add(YearGroupStudents(year_group_id=year_group_id, student_index=st.index))
        db.session.commit()
    return st


def create_groups(yg: YearGroup, amount: int) -> List[Group]:
    groups = [GroupFactory(year_group_id=yg.id) for _ in range(amount)]
    db.session.add_all(groups)
    db.session.commit()
    return groups


def create_group(data: dict, yg: YearGroup) -> Group:
    group = Group(**data, year_group_id=yg.id)
    db.session.add(group)
    db.session.commit()
    return group