code: def test_retry_failed_trial_callback(
      storage_mode: str, max_retry: Optional[int]
  ) -> None:
      heartbeat_interval = 1
      grace_period = 2
  
      with StorageSupplier(
          storage_mode,
          heartbeat_interval=heartbeat_interval,
          grace_period=grace_period,
          failed_trial_callback=RetryFailedTrialCallback(
              max_retry=max_retry
          ),
      ) as storage:
          assert is_heartbeat_enabled(storage)
          assert isinstance(storage, BaseHeartbeat)
  
          study = optuna.create_study(storage=storage)
  
          with pytest.warns(UserWarning):
              trial = study.ask()
          trial.suggest_float("_", -1, -1)
          trial.report(0.5, 1)
          storage.record_heartbeat(trial._trial_id)
          time.sleep(grace_period + 1)
  
          # Exceptions raised in spawned threads are caught by `_TestableThread`.
          with patch(
              "optuna.storages._heartbeat.Thread", _TestableThread
          ):
              study.optimize(lambda _: 1.0, n_trials=1)
  
          # Test the last trial to see if it was a retry of the first trial or not.
          # Test max_retry=None to see if trial is retried.
          # Test max_retry=0 to see if no trials are retried.
          # Test max_retry=1 to see if trial is retried.
          assert RetryFailedTrialCallback.retried_trial_number(
              study.trials[1]
          ) == (None if max_retry == 0 else 0)
          # Test inheritance of trial fields.
          if max_retry != 0:
              assert study.trials[0].params == study.trials[1].params
              assert (
                  study.trials[0].distributions
                  == study.trials[1].distributions
              )
              assert (
                  study.trials[0].user_attrs
                  == study.trials[1].user_attrs
              )
              # Only `intermediate_values` are not inherited.
              assert study.trials[1].intermediate_values == {}
  
imports: import itertools
  import multiprocessing
  import time
  from typing import Any
  from typing import Optional
  from unittest.mock import Mock
  from unittest.mock import patch
  import warnings
  import pytest
  import optuna
  from optuna import Study
  from optuna._callbacks import RetryFailedTrialCallback
  from optuna.storages import RDBStorage
  from optuna.storages._heartbeat import BaseHeartbeat
  from optuna.storages._heartbeat import is_heartbeat_enabled
  from optuna.testing.storages import STORAGE_MODES_HEARTBEAT
  from optuna.testing.storages import StorageSupplier
  from optuna.testing.threading import _TestableThread
  from optuna.trial import FrozenTrial
  from optuna.trial import TrialState
label_desc: Use image node_id 4 for calling a global function with example usage: test_retry_failed_trial_callback(storage_mode, max_retry) without return types