def test_fit_loop_save_and_restore_dataloaders(
    train_dataloader_factory, has_state, batches_before, batches_after, tmp_path
):
    """Test that the CheckpointConnector saves the state of stateful dataloaders."""

    class DummyModel(BoringModel):
        def __init__(self):
            super().__init__()
            self.seen_data = []

        def training_step(self, batch, batch_idx):
            self.seen_data.append(batch)
            print(batch)

        def train_dataloader(self):
            return train_dataloader_factory()

    trainer_kwargs = {
        "default_root_dir": tmp_path,
        "accelerator": "cpu",
        "enable_checkpointing": False,
        "enable_model_summary": False,
        "enable_progress_bar": False,
        "logger": False,
        "num_sanity_val_steps": 0,
    }

    # Train for 2 steps
    model = DummyModel()
    trainer = Trainer(**trainer_kwargs, max_steps=2)
    trainer.fit(model)
    assert model.seen_data == batches_before

    # Save a checkpoint
    trainer.save_checkpoint(tmp_path / "checkpoint.ckpt")
    checkpoint = torch.load(tmp_path / "checkpoint.ckpt", weights_only=True)
    if has_state:
        assert checkpoint["loops"]["fit_loop"]["state_dict"]["combined_loader"]
    else:
        assert "combined_loader" not in checkpoint["loops"]["fit_loop"]["state_dict"]

    # Restore training from step 2 and continue 2 more steps
    model = DummyModel()
    trainer = Trainer(**trainer_kwargs, max_steps=4)
    trainer.fit(model, ckpt_path=(tmp_path / "checkpoint.ckpt"))
    assert model.seen_data == batches_after
Test that the CheckpointConnector saves the state of stateful dataloaders.
tests/tests_pytorch/loops/test_loops.py