code: def test_trainer_fn_while_running(tmpdir, extra_params):
      class TestModel(BoringModel):
          def __init__(self, expected_fn, expected_stage):
              super().__init__()
              self.expected_fn = expected_fn
              self.expected_stage = expected_stage
              self.lr = 0.1
  
          def on_train_batch_start(self, *_):
              assert self.trainer.state.status == TrainerStatus.RUNNING
              assert self.trainer.state.fn == self.expected_fn
              assert self.trainer.training
  
          def on_sanity_check_start(self, *_):
              assert self.trainer.state.status == TrainerStatus.RUNNING
              assert self.trainer.state.fn == self.expected_fn
              assert self.trainer.sanity_checking
  
          def on_validation_batch_start(self, *_):
              assert self.trainer.state.status == TrainerStatus.RUNNING
              assert self.trainer.state.fn == self.expected_fn
              assert (
                  self.trainer.validating
                  or self.trainer.sanity_checking
              )
  
          def on_test_batch_start(self, *_):
              assert self.trainer.state.status == TrainerStatus.RUNNING
              assert self.trainer.state.fn == self.expected_fn
              assert self.trainer.testing
  
      trainer = Trainer(default_root_dir=tmpdir, **extra_params)
  
      model = TestModel(TrainerFn.FITTING, RunningStage.TRAINING)
      trainer.fit(model)
      assert trainer.state.finished
  
      model = TestModel(TrainerFn.VALIDATING, RunningStage.VALIDATING)
      trainer.validate(model)
      assert trainer.state.finished
  
      model = TestModel(TrainerFn.TESTING, RunningStage.TESTING)
      trainer.test(model)
      assert trainer.state.finished
  
imports: import pytest
  from lightning.pytorch import Callback, Trainer
  from lightning.pytorch.demos.boring_classes import BoringModel
  from lightning.pytorch.trainer.states import RunningStage, TrainerFn, TrainerState, TrainerStatus
label_desc: Use image node_id 2 for calling a global function with example usage: test_trainer_fn_while_running(tmpdir, extra_params) without return types