def test_count(session: Session) -> None:
study_1 = StudyModel(study_id=1, study_name="test-study-1")
study_2 = StudyModel(study_id=2, study_name="test-study-2")
session.add(TrialModel(study_id=study_1.study_id, state=TrialState.COMPLETE))
session.add(TrialModel(study_id=study_1.study_id, state=TrialState.RUNNING))
session.add(TrialModel(study_id=study_2.study_id, state=TrialState.RUNNING))
session.commit()
assert 3 == TrialModel.count(session)
assert 2 == TrialModel.count(session, study=study_1)
assert 1 == TrialModel.count(session, state=TrialState.COMPLETE)
def test_cascade_delete_on_study(session: Session) -> None:
study_id = 1
direction = StudyDirectionModel(direction=StudyDirection.MINIMIZE, objective=0)
study = StudyModel(study_id=study_id, study_name="test-study", directions=[direction])
study.trials.append(TrialModel(study_id=study.study_id, state=TrialState.COMPLETE))
study.trials.append(TrialModel(study_id=study.study_id, state=TrialState.RUNNING))
session.add(study)
session.commit()
assert 2 == TrialModel.count(session, study)
session.delete(study)
session.commit()
assert 0 == TrialModel.count(session, study)
