def test_trainer_setup_call(tmp_path, stage):
    """Test setup call gets the correct stage."""

    class CurrentModel(BoringModel):
        def setup(self, stage):
            self.stage = stage

    class CurrentCallback(Callback):
        def setup(self, trainer, model, stage):
            assert model is not None
            self.stage = stage

    model = CurrentModel()
    callback = CurrentCallback()
    trainer = Trainer(default_root_dir=tmp_path, max_epochs=1, enable_checkpointing=False, callbacks=[callback])

    if stage == "fit":
        trainer.fit(model)
    elif stage == "validate":
        trainer.validate(model)
    else:
        trainer.test(model)

    assert callback.stage == stage
    assert model.stage == stage
Test setup call gets the correct stage.
tests/tests_pytorch/trainer/test_trainer.py