def test_wandb_pickle(wandb_mock, tmp_path):
    """Verify that pickling trainer with wandb logger works.

    Wandb doesn't work well with pytest so we have to mock it out here.

    """

    class Experiment:
        id = "the_id"
        step = 0
        dir = "wandb"

        @property
        def name(self):
            return "the_run_name"

    wandb_mock.wandb_run = Experiment

    wandb_mock.run = None
    wandb_mock.init.return_value = Experiment()
    logger = WandbLogger(id="the_id", offline=True)

    trainer = Trainer(default_root_dir=tmp_path, max_epochs=1, logger=logger)
    # Access the experiment to ensure it's created
    assert trainer.logger.experiment, "missing experiment"
    assert trainer.log_dir == logger.save_dir
    pkl_bytes = pickle.dumps(trainer)
    trainer2 = pickle.loads(pkl_bytes)

    assert os.environ["WANDB_MODE"] == "dryrun"
    assert trainer2.logger.__class__.__name__ == WandbLogger.__name__
    assert trainer2.logger.experiment, "missing experiment"

    wandb_mock.init.assert_called()
    assert "id" in wandb_mock.init.call_args[1]
    assert wandb_mock.init.call_args[1]["id"] == "the_id"

    del os.environ["WANDB_MODE"]
Verify that pickling trainer with wandb logger works.

    Wandb doesn't work well with pytest so we have to mock it out here.

    
tests/tests_pytorch/loggers/test_wandb.py