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 _create_model(session: Session) -> StudyModel:
study = StudyModel(study_id=1, study_name="test-study")
dummy_study = StudyModel(study_id=2, study_name="dummy-study")
session.add(
StudyDirectionModel(study_id=study.study_id,
direction=StudyDirection.MINIMIZE,
objective=0))
session.add(
StudyDirectionModel(study_id=dummy_study.study_id,
direction=StudyDirection.MINIMIZE,
objective=0))
session.commit()
return study
def test_where_study_id(session: Session) -> None:
sample_study = StudyModel(study_id=1, study_name="test-study")
empty_study = StudyModel(study_id=2, study_name="test-study")
session.add(
StudySystemAttributeModel(
study_id=sample_study.study_id, key="sample-key", value_json="1"
)
)
assert 1 == len(StudySystemAttributeModel.where_study_id(sample_study.study_id, session))
assert 0 == len(StudySystemAttributeModel.where_study_id(empty_study.study_id, session))
# Check the case of unknown study_id.
assert 0 == len(StudySystemAttributeModel.where_study_id(-1, session))
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.system_attributes.append(
StudySystemAttributeModel(study_id=study_id,
key="sample-key1",
value_json="1"))
study.system_attributes.append(
StudySystemAttributeModel(study_id=study_id,
key="sample-key2",
value_json="2"))
session.add(study)
session.commit()
assert 2 == len(
StudySystemAttributeModel.where_study_id(study_id, session))
session.delete(study)
session.commit()
assert 0 == len(
StudySystemAttributeModel.where_study_id(study_id, session))
def test_cascade_delete_on_trial(session: Session) -> None:
trial_id = 1
direction = StudyDirectionModel(direction=StudyDirection.MINIMIZE,
objective=0)
study = StudyModel(study_id=1,
study_name="test-study",
directions=[direction])
trial = TrialModel(trial_id=trial_id,
study_id=study.study_id,
state=TrialState.COMPLETE)
trial.system_attributes.append(
TrialSystemAttributeModel(trial_id=trial_id,
key="sample-key1",
value_json="1"))
trial.system_attributes.append(
TrialSystemAttributeModel(trial_id=trial_id,
key="sample-key2",
value_json="2"))
study.trials.append(trial)
session.add(study)
session.commit()
assert 2 == len(
TrialSystemAttributeModel.where_trial_id(trial_id, session))
session.delete(trial)
session.commit()
assert 0 == len(
TrialSystemAttributeModel.where_trial_id(trial_id, session))
def test_count_past_trials(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")
trial_1_1 = TrialModel(study_id=study_1.study_id, state=TrialState.COMPLETE)
session.add(trial_1_1)
session.commit()
assert 0 == trial_1_1.count_past_trials(session)
trial_1_2 = TrialModel(study_id=study_1.study_id, state=TrialState.RUNNING)
session.add(trial_1_2)
session.commit()
assert 1 == trial_1_2.count_past_trials(session)
trial_2_1 = TrialModel(study_id=study_2.study_id, state=TrialState.RUNNING)
session.add(trial_2_1)
session.commit()
assert 0 == trial_2_1.count_past_trials(session)
def _create_model(session: Session) -> TrialModel:
direction = StudyDirectionModel(direction=StudyDirection.MINIMIZE, objective=0)
study = StudyModel(study_id=1, study_name="test-study", directions=[direction])
trial = TrialModel(trial_id=1, study_id=study.study_id, state=TrialState.COMPLETE)
session.add(study)
session.add(trial)
session.add(TrialHeartbeatModel(trial_id=trial.trial_id))
session.commit()
return trial
def test_find_by_study_and_key(session: Session) -> None:
study = StudyModel(study_id=1, study_name="test-study")
session.add(
StudySystemAttributeModel(study_id=study.study_id, key="sample-key", value_json="1")
)
session.commit()
attr = StudySystemAttributeModel.find_by_study_and_key(study, "sample-key", session)
assert attr is not None and "1" == attr.value_json
assert StudySystemAttributeModel.find_by_study_and_key(study, "not-found", session) is None
def test_all(session: Session) -> None:
study = StudyModel(study_id=1, study_name="test-study", direction=StudyDirection.MINIMIZE)
trial = TrialModel(trial_id=1, study_id=study.study_id, state=TrialState.COMPLETE)
session.add(
TrialSystemAttributeModel(trial_id=trial.trial_id, key="sample-key", value_json="1")
)
session.commit()
system_attributes = TrialSystemAttributeModel.all(session)
assert 1 == len(system_attributes)
assert "sample-key" == system_attributes[0].key
assert "1" == system_attributes[0].value_json
def test_find_by_trial_and_key(session: Session) -> None:
study = StudyModel(study_id=1, study_name="test-study")
trial = TrialModel(study_id=study.study_id)
session.add(
TrialSystemAttributeModel(trial_id=trial.trial_id, key="sample-key", value_json="1")
)
session.commit()
attr = TrialSystemAttributeModel.find_by_trial_and_key(trial, "sample-key", session)
assert attr is not None
assert "1" == attr.value_json
assert TrialSystemAttributeModel.find_by_trial_and_key(trial, "not-found", session) is None
def test_where_trial(session):
# type: (Session) -> None
study = StudyModel(study_id=1, study_name="test-study", direction=StudyDirection.MINIMIZE)
trial = TrialModel(trial_id=1, study_id=study.study_id, state=TrialState.COMPLETE)
session.add(
TrialUserAttributeModel(trial_id=trial.trial_id, key="sample-key", value_json="1")
)
session.commit()
user_attributes = TrialUserAttributeModel.where_trial(trial, session)
assert 1 == len(user_attributes)
assert "sample-key" == user_attributes[0].key
assert "1" == user_attributes[0].value_json
def test_cascade_delete_on_study(session: Session) -> None:
directions = [
StudyDirectionModel(study_id=1, direction=StudyDirection.MINIMIZE, objective=0),
StudyDirectionModel(study_id=1, direction=StudyDirection.MAXIMIZE, objective=1),
]
study = StudyModel(study_id=1, study_name="test-study", directions=directions)
session.add(study)
session.commit()
assert 2 == len(StudyDirectionModel.where_study_id(study.study_id, session))
session.delete(study)
session.commit()
assert 0 == len(StudyDirectionModel.where_study_id(study.study_id, session))
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)
def test_cascade_delete_on_study(session: Session) -> None:
study_id = 1
study = StudyModel(
study_id=study_id, study_name="test-study", direction=StudyDirection.MINIMIZE
)
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 == len(TrialModel.where_study(study, session))
session.delete(study)
session.commit()
assert 0 == len(TrialModel.where_study(study, session))
def _create_model(session: Session) -> TrialModel:
direction = StudyDirectionModel(direction=StudyDirection.MINIMIZE, objective=0)
study = StudyModel(study_id=1, study_name="test-study", directions=[direction])
trial = TrialModel(trial_id=1, study_id=study.study_id, state=TrialState.COMPLETE)
session.add(study)
session.add(trial)
session.add(
TrialIntermediateValueModel(
trial_id=trial.trial_id,
step=0,
intermediate_value=10,
intermediate_value_type=TrialIntermediateValueModel.TrialIntermediateValueType.FINITE, # noqa: E501
)
)
session.commit()
return trial
def test_all(session: Session) -> None:
direction = StudyDirectionModel(direction=StudyDirection.MINIMIZE,
objective=0)
study = StudyModel(study_id=1,
study_name="test-study",
directions=[direction])
trial = TrialModel(trial_id=1,
study_id=study.study_id,
state=TrialState.COMPLETE)
session.add(
TrialUserAttributeModel(trial_id=trial.trial_id,
key="sample-key",
value_json="1"))
session.commit()
user_attributes = TrialUserAttributeModel.all(session)
assert 1 == len(user_attributes)
assert "sample-key" == user_attributes[0].key
assert "1" == user_attributes[0].value_json
def test_cascade_delete_on_trial(session):
# type: (Session) -> None
trial_id = 1
study = StudyModel(study_id=1, study_name="test-study", direction=StudyDirection.MINIMIZE)
trial = TrialModel(trial_id=trial_id, study_id=study.study_id, state=TrialState.COMPLETE)
trial.user_attributes.append(
TrialUserAttributeModel(trial_id=trial_id, key="sample-key1", value_json="1")
)
trial.user_attributes.append(
TrialUserAttributeModel(trial_id=trial_id, key="sample-key2", value_json="2")
)
study.trials.append(trial)
session.add(study)
session.commit()
assert 2 == len(TrialUserAttributeModel.where_trial_id(trial_id, session))
session.delete(trial)
session.commit()
assert 0 == len(TrialUserAttributeModel.where_trial_id(trial_id, session))
