def test_record_heartbeat() -> None: heartbeat_interval = 1 n_trials = 2 sleep_sec = 2 def objective(trial: Trial) -> float: time.sleep(sleep_sec) return 1.0 with StorageSupplier("sqlite") as storage: assert isinstance(storage, RDBStorage) storage.heartbeat_interval = heartbeat_interval study = create_study(storage=storage) # Exceptions raised in spawned threads are caught by `_TestableThread`. with patch("optuna.study._optimize.Thread", _TestableThread): study.optimize(objective, n_trials=n_trials) trial_heartbeats = [] with _create_scoped_session(storage.scoped_session) as session: trial_ids = [trial._trial_id for trial in study.trials] for trial_id in trial_ids: heartbeat_model = TrialHeartbeatModel.where_trial_id( trial_id, session) assert heartbeat_model is not None trial_heartbeats.append(heartbeat_model.heartbeat) assert len(trial_heartbeats) == n_trials for i in range(n_trials - 1): assert (trial_heartbeats[i + 1] - trial_heartbeats[i]).seconds - sleep_sec <= 1
def test_create_scoped_session() -> None: storage = create_test_storage() # This object violates the unique constraint of version_info_id. v = VersionInfoModel(version_info_id=1, schema_version=1, library_version="0.0.1") with pytest.raises(IntegrityError): with _create_scoped_session(storage.scoped_session) as session: session.add(v)