def test_log_works_in_test_callback(tmp_path):
    """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=tmp_path, 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
Tests that log can be called within callback.
tests/tests_pytorch/trainer/logging_/test_eval_loop_logging.py