system-pri/backend/tests/fake_data.py

from typing import List

from .factory import ProjectSupervisorFactory, YearGroupProjectSupervisorsFactory, \
    StudentFactory, YearGroupStudentsFactory
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) -> List[ProjectSupervisorFactory]:
    ps = [ProjectSupervisorFactory() for _ in range(amount)]
    db.session.add_all(ps)
    db.session.commit()
    db.session.add_all(
        [YearGroupProjectSupervisorsFactory(limit_group=3, 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) -> Student:
    st = Student(**data)
    db.session.add(st)
    db.session.commit()
    db.session.add(YearGroupStudents(year_group_id=year_group_id, student_index=st.index))
    db.session.commit()
    return st