test: Move WandbLogger test to test_wandb

the trainer tests don't run with wandb installed, so we can't put it there

Author: ringohoffman

tests/tests_pytorch/loggers/test_wandb.py

@@ -24,7 +24,7 @@
 from lightning.pytorch.callbacks import ModelCheckpoint
 from lightning.pytorch.cli import LightningCLI
 from lightning.pytorch.demos.boring_classes import BoringModel
-from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch.loggers import TensorBoardLogger, WandbLogger
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
 from tests_pytorch.test_cli import _xfail_python_ge_3_11_9
 
@@ -133,6 +133,43 @@ def test_wandb_logger_init_before_spawn(wandb_mock):
     assert logger._experiment is not None
 
 
+def test_wandb_logger_experiment_called_first(wandb_mock, tmp_path):
+    wandb_experiment_called = False
+
+    def tensorboard_experiment_side_effect() -> mock.MagicMock:
+        nonlocal wandb_experiment_called
+        assert wandb_experiment_called
+        return mock.MagicMock()
+
+    def wandb_experiment_side_effect() -> mock.MagicMock:
+        nonlocal wandb_experiment_called
+        wandb_experiment_called = True
+        return mock.MagicMock()
+
+    with (
+        mock.patch.object(
+            TensorBoardLogger,
+            "experiment",
+            new_callable=lambda: mock.PropertyMock(side_effect=tensorboard_experiment_side_effect),
+        ),
+        mock.patch.object(
+            WandbLogger,
+            "experiment",
+            new_callable=lambda: mock.PropertyMock(side_effect=wandb_experiment_side_effect),
+        ),
+    ):
+        model = BoringModel()
+        trainer = Trainer(
+            default_root_dir=tmp_path,
+            log_every_n_steps=1,
+            limit_train_batches=0,
+            limit_val_batches=0,
+            max_steps=1,
+            logger=[TensorBoardLogger(tmp_path), WandbLogger(save_dir=tmp_path)],
+        )
+        trainer.fit(model)
+
+
 def test_wandb_pickle(wandb_mock, tmp_path):
     """Verify that pickling trainer with wandb logger works.
 

tests/tests_pytorch/trainer/test_trainer.py

@@ -49,7 +49,7 @@
     RandomIterableDataset,
     RandomIterableDatasetWithLen,
 )
-from lightning.pytorch.loggers import TensorBoardLogger, WandbLogger
+from lightning.pytorch.loggers import TensorBoardLogger
 from lightning.pytorch.overrides.distributed import UnrepeatedDistributedSampler, _IndexBatchSamplerWrapper
 from lightning.pytorch.strategies import DDPStrategy, SingleDeviceStrategy
 from lightning.pytorch.strategies.launchers import _MultiProcessingLauncher, _SubprocessScriptLauncher
@@ -1271,43 +1271,6 @@ def training_step(self, *args, **kwargs):
     log_metrics_mock.assert_has_calls(expected_calls)
 
 
-def test_wandb_logger_experiment_called_first(tmp_path):
-    wandb_experiment_called = False
-
-    def tensorboard_experiment_side_effect() -> mock.MagicMock:
-        nonlocal wandb_experiment_called
-        assert wandb_experiment_called
-        return mock.MagicMock()
-
-    def wandb_experiment_side_effect() -> mock.MagicMock:
-        nonlocal wandb_experiment_called
-        wandb_experiment_called = True
-        return mock.MagicMock()
-
-    with (
-        mock.patch.object(
-            TensorBoardLogger,
-            "experiment",
-            new_callable=lambda: mock.PropertyMock(side_effect=tensorboard_experiment_side_effect),
-        ),
-        mock.patch.object(
-            WandbLogger,
-            "experiment",
-            new_callable=lambda: mock.PropertyMock(side_effect=wandb_experiment_side_effect),
-        ),
-    ):
-        model = BoringModel()
-        trainer = Trainer(
-            default_root_dir=tmp_path,
-            log_every_n_steps=1,
-            limit_train_batches=0,
-            limit_val_batches=0,
-            max_steps=1,
-            logger=[TensorBoardLogger(tmp_path), WandbLogger(save_dir=tmp_path)],
-        )
-        trainer.fit(model)
-
-
 class TestLightningDataModule(LightningDataModule):
     def __init__(self, dataloaders):
         super().__init__()