code: def test_set_storage(tmp_path):
      """Test that setting the trial's storage context will update the tracked
      checkpoint paths."""
      original_storage = mock_storage_context()
      trial = Trial(
          "MockTrainable",
          stub=True,
          trial_id="abcd1234",
          storage=original_storage,
      )
  
      result_1 = _TrainingResult(
          checkpoint=Checkpoint.from_directory(
              original_storage.checkpoint_fs_path
          ),
          metrics={},
      )
      trial.on_checkpoint(result_1)
  
      result_2 = _TrainingResult(
          checkpoint=Checkpoint.from_directory(
              original_storage.checkpoint_fs_path
          ),
          metrics={},
      )
      trial.on_checkpoint(result_2)
  
      new_storage = StorageContext(
          storage_path=tmp_path / "new_storage_path",
          experiment_dir_name="new_name",
          trial_dir_name="new_trial",
      )
      trial.set_storage(new_storage)
  
      assert result_1.checkpoint.path.startswith(
          new_storage.trial_fs_path
      )
      assert result_2.checkpoint.path.startswith(
          new_storage.trial_fs_path
      )
  
file_path: ray/python/ray/tune/tests/test_trial.py
imports: import os
  import sys
  import pytest
  from ray.exceptions import RayActorError, RayTaskError
  from ray.train import Checkpoint
  from ray.train.constants import RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE
  from ray.train._internal.session import _TrainingResult
  from ray.train._internal.storage import StorageContext
  from ray.tune.experiment import Trial
  from ray.train.tests.util import mock_storage_context
label_desc: Use image node_id 4 for calling a global function with example usage: test_set_storage(tmp_path) without return types