def test_wandb_log_model(wandb_mock, tmp_path):
    """Test that the logger creates the folders and files in the right place."""
    wandb_mock.run = None
    model = BoringModel()

    # test log_model=True
    logger = WandbLogger(save_dir=tmp_path, log_model=True)
    logger.experiment.id = "1"
    logger.experiment.name = "run_name"
    trainer = Trainer(
        default_root_dir=tmp_path, logger=logger, max_epochs=2, limit_train_batches=3, limit_val_batches=3
    )
    trainer.fit(model)
    wandb_mock.init().log_artifact.assert_called_once()

    # test log_model='all'
    wandb_mock.init().log_artifact.reset_mock()
    wandb_mock.init.reset_mock()
    logger = WandbLogger(save_dir=tmp_path, log_model="all")
    logger.experiment.id = "1"
    logger.experiment.name = "run_name"
    trainer = Trainer(
        default_root_dir=tmp_path, logger=logger, max_epochs=2, limit_train_batches=3, limit_val_batches=3
    )
    trainer.fit(model)
    assert wandb_mock.init().log_artifact.call_count == 2

    # test log_model=False
    wandb_mock.init().log_artifact.reset_mock()
    wandb_mock.init.reset_mock()
    logger = WandbLogger(save_dir=tmp_path, log_model=False)
    logger.experiment.id = "1"
    logger.experiment.name = "run_name"
    trainer = Trainer(
        default_root_dir=tmp_path, logger=logger, max_epochs=2, limit_train_batches=3, limit_val_batches=3
    )
    trainer.fit(model)
    assert not wandb_mock.init().log_artifact.called

    # test correct metadata
    wandb_mock.init().log_artifact.reset_mock()
    wandb_mock.init.reset_mock()
    wandb_mock.Artifact.reset_mock()
    logger = WandbLogger(save_dir=tmp_path, log_model=True)
    logger.experiment.id = "1"
    logger.experiment.name = "run_name"
    trainer = Trainer(
        default_root_dir=tmp_path, logger=logger, max_epochs=2, limit_train_batches=3, limit_val_batches=3
    )
    trainer.fit(model)
    wandb_mock.Artifact.assert_called_once_with(
        name="model-1",
        type="model",
        metadata={
            "score": None,
            "original_filename": "epoch=1-step=6-v3.ckpt",
            "ModelCheckpoint": {
                "monitor": None,
                "mode": "min",
                "save_last": None,
                "save_top_k": 1,
                "save_weights_only": False,
                "_every_n_train_steps": 0,
            },
        },
    )

    # Test wandb custom artifact name
    wandb_mock.init().log_artifact.reset_mock()
    wandb_mock.init().reset_mock()
    wandb_mock.Artifact.reset_mock()
    logger = WandbLogger(save_dir=tmp_path, log_model=True, checkpoint_name="my-test-model")
    logger.experiment.id = "1"
    logger.experiment.name = "run_name"
    trainer = Trainer(
        default_root_dir=tmp_path, logger=logger, max_epochs=2, limit_train_batches=3, limit_val_batches=3
    )
    trainer.fit(model)
    wandb_mock.Artifact.assert_called_once_with(
        name="my-test-model",
        type="model",
        metadata={
            "score": None,
            "original_filename": "epoch=1-step=6-v4.ckpt",
            "ModelCheckpoint": {
                "monitor": None,
                "mode": "min",
                "save_last": None,
                "save_top_k": 1,
                "save_weights_only": False,
                "_every_n_train_steps": 0,
            },
        },
    )

    # Test wandb artifact with checkpoint_callback top_k logging latest
    wandb_mock.init().log_artifact.reset_mock()
    wandb_mock.init.reset_mock()
    wandb_mock.Artifact.reset_mock()
    logger = WandbLogger(save_dir=tmp_path, log_model=True)
    logger.experiment.id = "1"
    logger.experiment.name = "run_name"
    trainer = Trainer(
        default_root_dir=tmp_path,
        logger=logger,
        max_epochs=3,
        limit_train_batches=3,
        limit_val_batches=3,
        callbacks=[ModelCheckpoint(monitor="step", save_top_k=2)],
    )
    trainer.fit(model)
    wandb_mock.Artifact.assert_called_with(
        name="model-1",
        type="model",
        metadata={
            "score": 6,
            "original_filename": "epoch=1-step=6-v5.ckpt",
            "ModelCheckpoint": {
                "monitor": "step",
                "mode": "min",
                "save_last": None,
                "save_top_k": 2,
                "save_weights_only": False,
                "_every_n_train_steps": 0,
            },
        },
    )
    wandb_mock.init().log_artifact.assert_called_with(wandb_mock.Artifact(), aliases=["latest"])

    # Test wandb artifact with checkpoint_callback top_k logging latest and best
    wandb_mock.init().log_artifact.reset_mock()
    wandb_mock.init.reset_mock()
    wandb_mock.Artifact.reset_mock()
    logger = WandbLogger(save_dir=tmp_path, log_model=True)
    logger.experiment.id = "1"
    logger.experiment.name = "run_name"
    trainer = Trainer(
        default_root_dir=tmp_path,
        logger=logger,
        max_epochs=3,
        limit_train_batches=3,
        limit_val_batches=3,
        callbacks=[
            ModelCheckpoint(
                monitor="step",
            )
        ],
    )
    trainer.fit(model)
    wandb_mock.Artifact.assert_called_with(
        name="model-1",
        type="model",
        metadata={
            "score": 3,
            "original_filename": "epoch=0-step=3-v1.ckpt",
            "ModelCheckpoint": {
                "monitor": "step",
                "mode": "min",
                "save_last": None,
                "save_top_k": 1,
                "save_weights_only": False,
                "_every_n_train_steps": 0,
            },
        },
    )
    wandb_mock.init().log_artifact.assert_called_with(wandb_mock.Artifact(), aliases=["latest", "best"])

    # Test wandb artifact with two checkpoint_callbacks
    wandb_mock.init().log_artifact.reset_mock()
    wandb_mock.init.reset_mock()
    wandb_mock.Artifact.reset_mock()
    logger = WandbLogger(save_dir=tmp_path, log_model=True)
    logger.experiment.id = "1"
    logger.experiment.name = "run_name"
    trainer = Trainer(
        default_root_dir=tmp_path,
        logger=logger,
        max_epochs=3,
        limit_train_batches=3,
        limit_val_batches=3,
        callbacks=[
            ModelCheckpoint(monitor="epoch", save_top_k=2),
            ModelCheckpoint(monitor="step", save_top_k=2),
        ],
    )
    trainer.fit(model)
    for name, val, version in [("epoch", 0, 2), ("step", 3, 3)]:
        wandb_mock.Artifact.assert_any_call(
            name="model-1",
            type="model",
            metadata={
                "score": val,
                "original_filename": f"epoch=0-step=3-v{version}.ckpt",
                "ModelCheckpoint": {
                    "monitor": name,
                    "mode": "min",
                    "save_last": None,
                    "save_top_k": 2,
                    "save_weights_only": False,
                    "_every_n_train_steps": 0,
                },
            },
        )
        wandb_mock.init().log_artifact.assert_any_call(wandb_mock.Artifact(), aliases=["latest"])
Test that the logger creates the folders and files in the right place.
tests/tests_pytorch/loggers/test_wandb.py