code: def test_handle_preemption_error(
      trial: Trial, count_preemption_errors: bool, monkeypatch
  ):
      """Check that the Trial counts preemption errors correctly."""
      if count_preemption_errors:
          monkeypatch.setenv(RAY_TRAIN_COUNT_PREEMPTION_AS_FAILURE, "1")
  
      # Case 1: Directly raised (preemption) RayActorError
      class PreemptionRayActorError(RayActorError):
          def preempted(self) -> bool:
              return True
  
      err = PreemptionRayActorError()
      trial.handle_error(err)
      assert trial.num_failures == (1 if count_preemption_errors else 0)
  
      # Case 2: RayTaskError, where the cause is a (preemption) RayActorError
      wrapped_err = RayTaskError(
          function_name="test", traceback_str="traceback_str", cause=err
      )
      trial.handle_error(wrapped_err)
      assert trial.num_failures == (2 if count_preemption_errors else 0)
  
      # Case 3: Non-preemption error
      non_preempted_err = RayActorError()
      trial.handle_error(non_preempted_err)
      assert trial.num_failures == (3 if count_preemption_errors else 1)
  
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 2 for calling a global function with example usage: test_handle_preemption_error(trial, count_preemption_errors, monkeypatch) and returns: True