code: def test_log_works_in_test_callback(tmpdir):
      """Tests that log can be called within callback."""
  
      class TestCallback(callbacks.Callback):
          # helpers
          count = 0
          choices = [False, True]
  
          # used to compute expected values
          callback_funcs_called = collections.defaultdict(list)
          funcs_called_count = collections.defaultdict(int)
          funcs_attr = {}
  
          def make_logging(
              self, pl_module, func_name, on_steps, on_epochs, prob_bars
          ):
              original_func_name = func_name[:]
              self.funcs_called_count[original_func_name] += 1
  
              for idx, (on_step, on_epoch, prog_bar) in enumerate(
                  itertools.product(on_steps, on_epochs, prob_bars)
              ):
                  func_name = original_func_name[:]
                  custom_func_name = f"{idx}_{func_name}"
  
                  if not on_step and not on_epoch:
                      with pytest.raises(
                          MisconfigurationException,
                          match="is not useful",
                      ):
                          pl_module.log(
                              custom_func_name,
                              self.count,
                              on_step=on_step,
                              on_epoch=on_epoch,
                          )
                      continue
                  pl_module.log(
                      custom_func_name,
                      self.count,
                      on_step=on_step,
                      on_epoch=on_epoch,
                      prog_bar=prog_bar,
                  )
  
                  num_dl_ext = ""
                  dl_idx = pl_module.trainer._results.dataloader_idx
                  if dl_idx is not None:
                      num_dl_ext = f"/dataloader_idx_{dl_idx}"
                      func_name += num_dl_ext
  
                  # catch information for verification
                  self.callback_funcs_called[func_name].append(
                      [self.count]
                  )
                  self.funcs_attr[custom_func_name + num_dl_ext] = {
                      "on_step": on_step,
                      "on_epoch": on_epoch,
                      "prog_bar": prog_bar,
                      "func_name": func_name,
                  }
                  if on_step and on_epoch:
                      self.funcs_attr[
                          f"{custom_func_name}_step" + num_dl_ext
                      ] = {
                          "on_step": True,
                          "on_epoch": False,
                          "prog_bar": prog_bar,
                          "func_name": func_name,
                      }
  
                      self.funcs_attr[
                          f"{custom_func_name}_epoch" + num_dl_ext
                      ] = {
                          "on_step": False,
                          "on_epoch": True,
                          "prog_bar": prog_bar,
                          "func_name": func_name,
                      }
  
          def on_test_start(self, _, pl_module):
              self.make_logging(
                  pl_module,
                  "on_test_start",
                  on_steps=[False],
                  on_epochs=[True],
                  prob_bars=self.choices,
              )
  
          def on_test_epoch_start(self, _, pl_module):
              self.make_logging(
                  pl_module,
                  "on_test_epoch_start",
                  on_steps=[False],
                  on_epochs=[True],
                  prob_bars=self.choices,
              )
  
          def on_test_batch_end(self, _, pl_module, *__):
              self.make_logging(
                  pl_module,
                  "on_test_batch_end",
                  on_steps=self.choices,
                  on_epochs=self.choices,
                  prob_bars=self.choices,
              )
  
          def on_test_epoch_end(self, _, pl_module):
              self.make_logging(
                  pl_module,
                  "on_test_epoch_end",
                  on_steps=[False],
                  on_epochs=[True],
                  prob_bars=self.choices,
              )
  
      num_dataloaders = 2
  
      class TestModel(BoringModel):
          seen_losses = {i: [] for i in range(num_dataloaders)}
  
          def test_step(self, batch, batch_idx, dataloader_idx=0):
              loss = super().test_step(batch, batch_idx)["y"]
              self.log("test_loss", loss)
              self.seen_losses[dataloader_idx].append(loss)
  
          def test_dataloader(self):
              return [
                  torch.utils.data.DataLoader(RandomDataset(32, 64))
                  for _ in range(num_dataloaders)
              ]
  
      model = TestModel()
      cb = TestCallback()
      trainer = Trainer(
          default_root_dir=tmpdir,
          limit_test_batches=2,
          num_sanity_val_steps=0,
          max_epochs=2,
          callbacks=[cb],
      )
      trainer.test(model)
  
      assert cb.funcs_called_count["on_test_start"] == 1
      assert cb.funcs_called_count["on_test_epoch_start"] == 1
      assert cb.funcs_called_count["on_test_batch_end"] == 4
      assert cb.funcs_called_count["on_test_epoch_end"] == 1
  
      callback_metrics = trainer.callback_metrics
      for func_name in cb.callback_funcs_called:
          for key in callback_metrics:
              if func_name in key:
                  break
          else:
              pytest.fail(f"{func_name}, {list(callback_metrics)}")
  
      def get_expected(on_epoch, values):
          reduction = np.mean if on_epoch else np.max
          return reduction(values)
  
      # Make sure the func_name output equals the average from all logged values when on_epoch true
      for dl_idx in range(num_dataloaders):
          key = f"test_loss/dataloader_idx_{dl_idx}"
          assert key in trainer.callback_metrics
          assert torch.stack(
              model.seen_losses[dl_idx]
          ).mean() == trainer.callback_metrics.pop(key)
  
      for func_name, output_value in trainer.callback_metrics.items():
          output_value = output_value.item()
          func_attr = cb.funcs_attr[func_name]
          original_values = cb.callback_funcs_called[
              func_attr["func_name"]
          ]
          expected_output = get_expected(
              func_attr["on_epoch"], original_values
          )
          assert output_value == expected_output
  
      for fx, attrs in cb.funcs_attr.items():
          should_include = (
              attrs["prog_bar"] and attrs["on_step"] ^ attrs["on_epoch"]
          )
          is_included = fx in trainer.progress_bar_metrics
          assert is_included if should_include else not is_included
  
imports: import collections
  import itertools
  import os
  from contextlib import redirect_stdout
  from io import StringIO
  from unittest import mock
  from unittest.mock import ANY, call
  import numpy
  import pytest
  import torch
  from lightning.fabric.utilities.imports import _PYTHON_GREATER_EQUAL_3_8_0
  from lightning.pytorch import Trainer, callbacks
  from lightning.pytorch.callbacks.progress.rich_progress import _RICH_AVAILABLE
  from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset
  from lightning.pytorch.loggers import TensorBoardLogger
  from lightning.pytorch.loops import _EvaluationLoop
  from lightning.pytorch.trainer.states import RunningStage
  from lightning.pytorch.utilities.exceptions import MisconfigurationException
  from torch import Tensor
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 9 for calling a global function with example usage: test_log_works_in_test_callback(tmpdir) and returns: reduction, unknown