from datetime import date from math import ceil from random import randint import click from flask_security.utils import hash_password from advlabdb import create_app, data_dir, user_datastore from advlabdb.exceptions import DatabaseException from advlabdb.models import ( Admin, Appointment, Assistant, Experiment, Group, GroupExperiment, Part, PartStudent, Program, Semester, SemesterExperiment, Student, db, ) from cli.test.generate_test_db.assistant_names import assistant_names from cli.test.generate_test_db.experiment_titles import experiment_titles from cli.test.generate_test_db.student_names import student_names def db_add(obj): db.session.add(obj) return obj def _generate_test_db(): db_file = data_dir / "db/advlabdb.sqlite" if db_file.is_file(): click.echo( click.style( f"Generating a test database is not allowed because a database does already exist at {db_file}.", fg="red", ) ) return app = create_app(create_for_server=False) with app.app_context(), db.session.begin(): db.create_all() bs_prog = db_add(Program(label="BS")) ms_prog = db_add(Program(label="MS")) be_prog = db_add(Program(label="BE")) programs = (bs_prog, ms_prog, be_prog) ss = db_add(Semester(label="SS", year=23)) ws = db_add(Semester(label="WS", year=23)) semesters = (ss, ws) bs_1_ss_part = db_add(Part(semester=ss, program=bs_prog, number=1)) bs_2_ss_part = db_add(Part(semester=ss, program=bs_prog, number=2)) ms_1_ss_part = db_add(Part(semester=ss, program=ms_prog, number=1)) ms_2_ss_part = db_add(Part(semester=ss, program=ms_prog, number=2)) be_1_ss_part = db_add(Part(semester=ss, program=be_prog, number=1)) be_2_ss_part = db_add(Part(semester=ss, program=be_prog, number=2)) ss_parts = (bs_1_ss_part, bs_2_ss_part, ms_1_ss_part, ms_2_ss_part, be_1_ss_part, be_2_ss_part) db_add(Part(semester=ws, program=bs_prog, number=1)) bs_2_ws_part = db_add(Part(semester=ws, program=bs_prog, number=2)) ms_1_ws_part = db_add(Part(semester=ws, program=ms_prog, number=1)) ms_2_ws_part = db_add(Part(semester=ws, program=ms_prog, number=2)) db_add(Part(semester=ws, program=be_prog, number=1)) be_2_ws_part = db_add(Part(semester=ws, program=be_prog, number=2)) semester_program_part_students = { ss: {program: [] for program in programs}, ws: {program: [] for program in programs}, } def add_part_student(student, part): part_student = db_add(PartStudent(student=student, part=part)) semester_program_part_students[part.semester][part.program].append(part_student) for ind, name in enumerate(student_names): contact_email = f"{name[0]}-{name[1]}@private.de".lower() if ind % 2 == 0 else None student = db_add( Student( student_number=ind, first_name=name[0], last_name=name[1], uni_email=f"{name[0]}-{name[1]}@uni.de".lower(), contact_email=contact_email, ) ) part = ss_parts[ind % len(ss_parts)] add_part_student(student, part) if ind % 10 < 9: # Transfer to the next part in the second semester if part == bs_1_ss_part: add_part_student(student, bs_2_ws_part) elif part == bs_2_ss_part: add_part_student(student, ms_1_ws_part) elif part == ms_1_ss_part: add_part_student(student, ms_2_ws_part) elif part == be_1_ss_part: add_part_student(student, be_2_ws_part) program_groups = {program: [] for program in programs} for program_part_students in semester_program_part_students.values(): for program, part_students in program_part_students.items(): if len(part_students) % 2 == 1: # Add the first 3 students into a special group for an uneven number of students group_part_students = [] for _ in range(3): if len(part_students) == 0: break group_part_students.append(part_students.pop(0)) group = db_add(Group.customInit(group_part_students)) program_groups[program].append(group) while len(part_students) >= 2: # Add the rest of the students into groups of 2 group = db_add(Group.customInit([part_students.pop(0) for i in range(2)])) program_groups[program].append(group) program_semester_experiments = {program: [] for program in programs} all_semester_experiments = [] for ind, title in enumerate(experiment_titles): program = programs[ind % len(programs)] experiment = db_add( Experiment( number=ind + 1, program=program, title=title, building="Not assigned", room="Not assigned", duration_in_days=2, ) ) for semester in semesters: semester_experiment = db_add(SemesterExperiment(experiment=experiment, semester=semester)) program_semester_experiments[program].append(semester_experiment) all_semester_experiments.append(semester_experiment) all_group_experiments = [] for program in programs: semester_experiments = program_semester_experiments[program] num_semester_experiments = len(semester_experiments) for ind, group in enumerate(program_groups[program]): num_tries = 0 while True: try: semester_experiment = semester_experiments[ind % num_semester_experiments] group_experiment = db_add(GroupExperiment(semester_experiment=semester_experiment, group=group)) except DatabaseException as ex: # Catch an error when a student becomes the same experiment for the second time. # Try the next experiment! ind += 1 num_tries += 1 if num_tries == num_semester_experiments: raise ex else: break all_group_experiments.append(group_experiment) adminRole = user_datastore.create_role(name="admin") assistantRole = user_datastore.create_role(name="assistant") admin_names = ( ("Aileen", "Grant"), ("Ben", "Huber"), ) for name in admin_names: user = user_datastore.create_user( email=f"{name[0]}-{name[1]}@uni.de".lower(), password=hash_password("admin"), roles=[adminRole], first_name=name[0], last_name=name[1], active_semester=ws, ) db_add(Admin(user=user)) num_assistants = len(assistant_names) num_assistant_experiments = ceil(len(all_semester_experiments) / num_assistants) for name in assistant_names: user = user_datastore.create_user( email=f"{name[0]}-{name[1]}@uni.de".lower(), password=hash_password("assistant"), roles=[assistantRole], first_name=name[0], last_name=name[1], mobile_phone_number="01" + "".join(str(randint(0, 9)) for i in range(10)), # nosec: B311 active_semester=ws, ) semester_experiments = [] for _ in range(num_assistant_experiments): if len(all_semester_experiments) == 0: break semester_experiments.append(all_semester_experiments.pop(0)) db_add(Assistant(user=user, semester_experiments=semester_experiments)) for group_experiment in all_group_experiments: semester_experiment = group_experiment.semester_experiment special = False semester = semester_experiment.semester if semester.label == "SS": month = randint(3, 8) # nosec: B311 if month <= 4: special = True else: month = randint(9, 12) # nosec: B311 if month <= 10: special = True year = 2000 + semester.year day = randint(1, 28) # nosec: B311 for appointment_date in (date(year, month, day), date(year, month, day + 1)): db_add( Appointment( date=appointment_date, special=special, group_experiment=group_experiment, assistant=semester_experiment.assistants[0], ) )