# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re
import shutil
from pathlib import Path

import pytest
import pytorch_lightning as pl
from omegaconf.errors import OmegaConfBaseException

from nemo.constants import NEMO_ENV_VARNAME_VERSION
from nemo.utils.exp_manager import (
    CheckpointMisconfigurationError,
    LoggerMisconfigurationError,
    NotFoundError,
    exp_manager,
)


@pytest.fixture
def cleanup_local_folder():
    # Asserts in fixture are not recommended, but I'd rather stop users from deleting expensive training runs
    assert not Path("./lightning_logs").exists()
    assert not Path("./NeMo_experiments").exists()
    assert not Path("./nemo_experiments").exists()

    yield

    if Path("./lightning_logs").exists():
        shutil.rmtree('./lightning_logs')
    if Path("./NeMo_experiments").exists():
        shutil.rmtree('./NeMo_experiments')
    if Path("./nemo_experiments").exists():
        shutil.rmtree('./nemo_experiments')


class TestExpManager:
    @pytest.mark.unit
    def test_omegaconf(self):
        """Ensure omegaconf raises an error when an unexcepted argument is passed"""
        with pytest.raises(OmegaConfBaseException):
            exp_manager(None, {"unused": 1})

    @pytest.mark.unit
    def test_trainer_loggers(self, cleanup_local_folder, tmp_path):
        """ Test that a trainer with logger errors out with a number of arguments. Test that it works with
        create_tensorboard_logger set to False
        """
        test_trainer = pl.Trainer()  # Should create logger and modelcheckpoint

        with pytest.raises(LoggerMisconfigurationError):  # Fails because exp_manager defaults to trainer
            exp_manager(test_trainer, {"exp_dir": str(tmp_path)})
        with pytest.raises(LoggerMisconfigurationError):  # Fails because exp_manager defaults to trainer
            exp_manager(test_trainer, {"explicit_log_dir": str(tmp_path)})
        with pytest.raises(LoggerMisconfigurationError):  # Fails because exp_manager defaults to trainer
            exp_manager(test_trainer, {"resume_if_exists": True})

        # Check that exp_manager uses trainer.logger, it's exp_dir, name, and version
        log_dir = exp_manager(test_trainer, {"create_tensorboard_logger": False, "create_checkpoint_callback": False})
        assert log_dir.resolve() == Path("./lightning_logs/version_0").resolve()
        assert Path("./lightning_logs").exists()
        assert Path("./lightning_logs/version_0").exists()

        # Check that a trainer without a logger gets a logger attached to it
        test_trainer = pl.Trainer(logger=False)
        log_dir = exp_manager(
            test_trainer,
            {"create_tensorboard_logger": True, "create_checkpoint_callback": False, "exp_dir": str(tmp_path)},
        )
        assert isinstance(test_trainer.logger, pl.loggers.TensorBoardLogger)

        test_trainer = pl.Trainer(logger=False)
        # Check that a create_wandb_logger=True errors out unless wandb_logger_kwargs is passed.
        with pytest.raises(ValueError):
            log_dir = exp_manager(
                test_trainer,
                {
                    "create_tensorboard_logger": False,
                    "create_checkpoint_callback": False,
                    "exp_dir": str(tmp_path),
                    "create_wandb_logger": True,
                },
            )
        # Check that a WandbLogger is attached to logger if create_wandb_logger=True and wandb_logger_kwargs has name
        # and project
        log_dir = exp_manager(
            test_trainer,
            {
                "create_tensorboard_logger": False,
                "create_checkpoint_callback": False,
                "exp_dir": str(tmp_path),
                "create_wandb_logger": True,
                "wandb_logger_kwargs": {"name": "", "project": ""},
            },
        )
        assert isinstance(test_trainer.logger, pl.loggers.WandbLogger)

    @pytest.mark.unit
    def test_checkpoint_configurations(self, cleanup_local_folder):
        """ Test that trainer creating modelcheckpoint and asking exp_manager to do it too results in errors, but
        is error free if only one is asked to do so.
        """
        disable_tb_logger = {"create_tensorboard_logger": False}
        test_trainer = pl.Trainer()  # Should create logger and modelcheckpoint
        with pytest.raises(CheckpointMisconfigurationError):  # Fails because both try to create modelcheckpoint
            exp_manager(test_trainer, disable_tb_logger)

        # Should succeed without error
        exp_manager(test_trainer, {"create_checkpoint_callback": False, "create_tensorboard_logger": False})

        test_trainer_2 = pl.Trainer(checkpoint_callback=False)
        exp_manager(test_trainer_2, disable_tb_logger)  # Should succeed without error

    @pytest.mark.unit
    def test_default_log_dir(self, cleanup_local_folder):
        """Check the default of ./nemo_experiments/default/datetime works as intended"""
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)

        log_dir = exp_manager(test_trainer, {"create_tensorboard_logger": False, "create_checkpoint_callback": False})
        assert (log_dir / "..").resolve() == Path("./nemo_experiments/default/").resolve()
        assert Path("./nemo_experiments").exists()
        assert Path("./nemo_experiments/default/").exists()
        sub_dirs = [x for x in Path("./nemo_experiments/default/").iterdir() if x.is_dir()]
        assert len(sub_dirs) == 1
        assert re.match(r"[0-9]{4}-[0-9]{2}-[0-9]{2}_[0-9]{2}-[0-9]{2}-[0-9]{2}", sub_dirs[0].name)

    @pytest.mark.unit
    def test_log_dir_overrides(self, monkeypatch, tmp_path):
        """Check a variety of trainer options with exp_manager"""
        # Checks that explicit_log_dir ignores exp_dir, name, and version
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {"explicit_log_dir": str(tmp_path / "test_log_dir_overrides")})
        assert log_dir.resolve() == (tmp_path / "test_log_dir_overrides").resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / "test_log_dir_overrides").exists()

        # Checks that exp_manager uses exp_dir, default name, and explicit version
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {"exp_dir": str(tmp_path / "test_no_name"), "version": 957})
        assert log_dir.resolve() == (tmp_path / "test_no_name" / "default" / "957").resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / "test_no_name" / "default" / "957").exists()

        monkeypatch.delenv(NEMO_ENV_VARNAME_VERSION)
        # Checks that use_datetime_version False toggle works
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {"exp_dir": str(tmp_path / "test_no_name"), "use_datetime_version": False})
        assert log_dir.resolve() == (tmp_path / "test_no_name" / "default" / "version_0").resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / "test_no_name" / "default" / "version_0").exists()

        monkeypatch.delenv(NEMO_ENV_VARNAME_VERSION)
        # Checks that use_datetime_version False toggle works and version increments
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {"exp_dir": str(tmp_path / "test_no_name"), "use_datetime_version": False})
        assert log_dir.resolve() == (tmp_path / "test_no_name" / "default" / "version_1").resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / "test_no_name" / "default" / "version_1").exists()

    @pytest.mark.unit
    def test_resume(self, tmp_path):
        """ Tests the resume capabilities of exp_manager"""
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)

        # Error because explicit_log_dir does not exist
        with pytest.raises(NotFoundError):
            exp_manager(
                test_trainer,
                {
                    "exp_dir": str(tmp_path / "test_resume"),
                    "resume_if_exists": True,
                    "explicit_log_dir": "Does_not_exist",
                },
            )

        # Error because checkpoints folder does not exist
        with pytest.raises(NotFoundError):
            exp_manager(test_trainer, {"resume_if_exists": True, "exp_dir": str(tmp_path / "test_resume")})

        # No error because we tell exp_manager to ignore notfounderror
        exp_manager(
            test_trainer,
            {
                "resume_if_exists": True,
                "exp_dir": str(tmp_path / "test_resume_2"),
                "resume_ignore_no_checkpoint": True,
            },
        )

        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints").mkdir(parents=True)
        # Error because checkpoints do not exist in folder
        with pytest.raises(NotFoundError):
            exp_manager(
                test_trainer,
                {
                    "resume_if_exists": True,
                    "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0"),
                },
            )

        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--end.ckpt").touch()
        # Error because *end.ckpt is in folder indicating that training has already finished
        with pytest.raises(ValueError):
            exp_manager(
                test_trainer,
                {
                    "resume_if_exists": True,
                    "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0"),
                },
            )

        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--end.ckpt").unlink()
        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--last.ckpt").touch()
        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel2--last.ckpt").touch()
        # Error because multiple *last.ckpt is in folder. If more than one, don't know which to restore
        with pytest.raises(ValueError):
            exp_manager(
                test_trainer,
                {
                    "resume_if_exists": True,
                    "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0"),
                },
            )

        # Finally succeed
        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel2--last.ckpt").unlink()
        log_dir = exp_manager(
            test_trainer,
            {"resume_if_exists": True, "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0")},
        )
        checkpoint = Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--last.ckpt")
        assert Path(test_trainer.resume_from_checkpoint).resolve() == checkpoint.resolve()

        # Succeed again and make sure that run_0 exists and previous log files were moved
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        exp_manager(test_trainer, {"resume_if_exists": True, "explicit_log_dir": str(log_dir)})
        checkpoint = Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--last.ckpt")
        assert Path(test_trainer.resume_from_checkpoint).resolve() == checkpoint.resolve()
        prev_run_dir = Path(tmp_path / "test_resume" / "default" / "version_0" / "run_0")
        assert prev_run_dir.exists()
        prev_log = Path(tmp_path / "test_resume" / "default" / "version_0" / "run_0" / "lightning_logs.txt")
        assert prev_log.exists()

acb4b1d054fc7a3aaaf32cafd5682d291a7741fa
<mask token>


class TestExpManager:

    @pytest.mark.unit
    def test_omegaconf(self):
        """Ensure omegaconf raises an error when an unexcepted argument is passed"""
        with pytest.raises(OmegaConfBaseException):
            exp_manager(None, {'unused': 1})
    <mask token>
    <mask token>
    <mask token>
    <mask token>
    <mask token>

<mask token>


class TestExpManager:

    @pytest.mark.unit
    def test_omegaconf(self):
        """Ensure omegaconf raises an error when an unexcepted argument is passed"""
        with pytest.raises(OmegaConfBaseException):
            exp_manager(None, {'unused': 1})

    @pytest.mark.unit
    def test_trainer_loggers(self, cleanup_local_folder, tmp_path):
        """ Test that a trainer with logger errors out with a number of arguments. Test that it works with
        create_tensorboard_logger set to False
        """
        test_trainer = pl.Trainer()
        with pytest.raises(LoggerMisconfigurationError):
            exp_manager(test_trainer, {'exp_dir': str(tmp_path)})
        with pytest.raises(LoggerMisconfigurationError):
            exp_manager(test_trainer, {'explicit_log_dir': str(tmp_path)})
        with pytest.raises(LoggerMisconfigurationError):
            exp_manager(test_trainer, {'resume_if_exists': True})
        log_dir = exp_manager(test_trainer, {'create_tensorboard_logger': 
            False, 'create_checkpoint_callback': False})
        assert log_dir.resolve() == Path('./lightning_logs/version_0').resolve(
            )
        assert Path('./lightning_logs').exists()
        assert Path('./lightning_logs/version_0').exists()
        test_trainer = pl.Trainer(logger=False)
        log_dir = exp_manager(test_trainer, {'create_tensorboard_logger': 
            True, 'create_checkpoint_callback': False, 'exp_dir': str(
            tmp_path)})
        assert isinstance(test_trainer.logger, pl.loggers.TensorBoardLogger)
        test_trainer = pl.Trainer(logger=False)
        with pytest.raises(ValueError):
            log_dir = exp_manager(test_trainer, {
                'create_tensorboard_logger': False,
                'create_checkpoint_callback': False, 'exp_dir': str(
                tmp_path), 'create_wandb_logger': True})
        log_dir = exp_manager(test_trainer, {'create_tensorboard_logger': 
            False, 'create_checkpoint_callback': False, 'exp_dir': str(
            tmp_path), 'create_wandb_logger': True, 'wandb_logger_kwargs':
            {'name': '', 'project': ''}})
        assert isinstance(test_trainer.logger, pl.loggers.WandbLogger)

    @pytest.mark.unit
    def test_checkpoint_configurations(self, cleanup_local_folder):
        """ Test that trainer creating modelcheckpoint and asking exp_manager to do it too results in errors, but
        is error free if only one is asked to do so.
        """
        disable_tb_logger = {'create_tensorboard_logger': False}
        test_trainer = pl.Trainer()
        with pytest.raises(CheckpointMisconfigurationError):
            exp_manager(test_trainer, disable_tb_logger)
        exp_manager(test_trainer, {'create_checkpoint_callback': False,
            'create_tensorboard_logger': False})
        test_trainer_2 = pl.Trainer(checkpoint_callback=False)
        exp_manager(test_trainer_2, disable_tb_logger)

    @pytest.mark.unit
    def test_default_log_dir(self, cleanup_local_folder):
        """Check the default of ./nemo_experiments/default/datetime works as intended"""
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {'create_tensorboard_logger': 
            False, 'create_checkpoint_callback': False})
        assert (log_dir / '..').resolve() == Path('./nemo_experiments/default/'
            ).resolve()
        assert Path('./nemo_experiments').exists()
        assert Path('./nemo_experiments/default/').exists()
        sub_dirs = [x for x in Path('./nemo_experiments/default/').iterdir(
            ) if x.is_dir()]
        assert len(sub_dirs) == 1
        assert re.match('[0-9]{4}-[0-9]{2}-[0-9]{2}_[0-9]{2}-[0-9]{2}-[0-9]{2}'
            , sub_dirs[0].name)
    <mask token>

    @pytest.mark.unit
    def test_resume(self, tmp_path):
        """ Tests the resume capabilities of exp_manager"""
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        with pytest.raises(NotFoundError):
            exp_manager(test_trainer, {'exp_dir': str(tmp_path /
                'test_resume'), 'resume_if_exists': True,
                'explicit_log_dir': 'Does_not_exist'})
        with pytest.raises(NotFoundError):
            exp_manager(test_trainer, {'resume_if_exists': True, 'exp_dir':
                str(tmp_path / 'test_resume')})
        exp_manager(test_trainer, {'resume_if_exists': True, 'exp_dir': str
            (tmp_path / 'test_resume_2'), 'resume_ignore_no_checkpoint': True})
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' / 'checkpoints'
            ).mkdir(parents=True)
        with pytest.raises(NotFoundError):
            exp_manager(test_trainer, {'resume_if_exists': True,
                'explicit_log_dir': str(tmp_path / 'test_resume' /
                'default' / 'version_0')})
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel--end.ckpt').touch()
        with pytest.raises(ValueError):
            exp_manager(test_trainer, {'resume_if_exists': True,
                'explicit_log_dir': str(tmp_path / 'test_resume' /
                'default' / 'version_0')})
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel--end.ckpt').unlink()
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel--last.ckpt').touch()
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel2--last.ckpt').touch()
        with pytest.raises(ValueError):
            exp_manager(test_trainer, {'resume_if_exists': True,
                'explicit_log_dir': str(tmp_path / 'test_resume' /
                'default' / 'version_0')})
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel2--last.ckpt').unlink()
        log_dir = exp_manager(test_trainer, {'resume_if_exists': True,
            'explicit_log_dir': str(tmp_path / 'test_resume' / 'default' /
            'version_0')})
        checkpoint = Path(tmp_path / 'test_resume' / 'default' /
            'version_0' / 'checkpoints' / 'mymodel--last.ckpt')
        assert Path(test_trainer.resume_from_checkpoint).resolve(
            ) == checkpoint.resolve()
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        exp_manager(test_trainer, {'resume_if_exists': True,
            'explicit_log_dir': str(log_dir)})
        checkpoint = Path(tmp_path / 'test_resume' / 'default' /
            'version_0' / 'checkpoints' / 'mymodel--last.ckpt')
        assert Path(test_trainer.resume_from_checkpoint).resolve(
            ) == checkpoint.resolve()
        prev_run_dir = Path(tmp_path / 'test_resume' / 'default' /
            'version_0' / 'run_0')
        assert prev_run_dir.exists()
        prev_log = Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'run_0' / 'lightning_logs.txt')
        assert prev_log.exists()

<mask token>


@pytest.fixture
def cleanup_local_folder():
    assert not Path('./lightning_logs').exists()
    assert not Path('./NeMo_experiments').exists()
    assert not Path('./nemo_experiments').exists()
    yield
    if Path('./lightning_logs').exists():
        shutil.rmtree('./lightning_logs')
    if Path('./NeMo_experiments').exists():
        shutil.rmtree('./NeMo_experiments')
    if Path('./nemo_experiments').exists():
        shutil.rmtree('./nemo_experiments')


class TestExpManager:

    @pytest.mark.unit
    def test_omegaconf(self):
        """Ensure omegaconf raises an error when an unexcepted argument is passed"""
        with pytest.raises(OmegaConfBaseException):
            exp_manager(None, {'unused': 1})

    @pytest.mark.unit
    def test_trainer_loggers(self, cleanup_local_folder, tmp_path):
        """ Test that a trainer with logger errors out with a number of arguments. Test that it works with
        create_tensorboard_logger set to False
        """
        test_trainer = pl.Trainer()
        with pytest.raises(LoggerMisconfigurationError):
            exp_manager(test_trainer, {'exp_dir': str(tmp_path)})
        with pytest.raises(LoggerMisconfigurationError):
            exp_manager(test_trainer, {'explicit_log_dir': str(tmp_path)})
        with pytest.raises(LoggerMisconfigurationError):
            exp_manager(test_trainer, {'resume_if_exists': True})
        log_dir = exp_manager(test_trainer, {'create_tensorboard_logger': 
            False, 'create_checkpoint_callback': False})
        assert log_dir.resolve() == Path('./lightning_logs/version_0').resolve(
            )
        assert Path('./lightning_logs').exists()
        assert Path('./lightning_logs/version_0').exists()
        test_trainer = pl.Trainer(logger=False)
        log_dir = exp_manager(test_trainer, {'create_tensorboard_logger': 
            True, 'create_checkpoint_callback': False, 'exp_dir': str(
            tmp_path)})
        assert isinstance(test_trainer.logger, pl.loggers.TensorBoardLogger)
        test_trainer = pl.Trainer(logger=False)
        with pytest.raises(ValueError):
            log_dir = exp_manager(test_trainer, {
                'create_tensorboard_logger': False,
                'create_checkpoint_callback': False, 'exp_dir': str(
                tmp_path), 'create_wandb_logger': True})
        log_dir = exp_manager(test_trainer, {'create_tensorboard_logger': 
            False, 'create_checkpoint_callback': False, 'exp_dir': str(
            tmp_path), 'create_wandb_logger': True, 'wandb_logger_kwargs':
            {'name': '', 'project': ''}})
        assert isinstance(test_trainer.logger, pl.loggers.WandbLogger)

    @pytest.mark.unit
    def test_checkpoint_configurations(self, cleanup_local_folder):
        """ Test that trainer creating modelcheckpoint and asking exp_manager to do it too results in errors, but
        is error free if only one is asked to do so.
        """
        disable_tb_logger = {'create_tensorboard_logger': False}
        test_trainer = pl.Trainer()
        with pytest.raises(CheckpointMisconfigurationError):
            exp_manager(test_trainer, disable_tb_logger)
        exp_manager(test_trainer, {'create_checkpoint_callback': False,
            'create_tensorboard_logger': False})
        test_trainer_2 = pl.Trainer(checkpoint_callback=False)
        exp_manager(test_trainer_2, disable_tb_logger)

    @pytest.mark.unit
    def test_default_log_dir(self, cleanup_local_folder):
        """Check the default of ./nemo_experiments/default/datetime works as intended"""
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {'create_tensorboard_logger': 
            False, 'create_checkpoint_callback': False})
        assert (log_dir / '..').resolve() == Path('./nemo_experiments/default/'
            ).resolve()
        assert Path('./nemo_experiments').exists()
        assert Path('./nemo_experiments/default/').exists()
        sub_dirs = [x for x in Path('./nemo_experiments/default/').iterdir(
            ) if x.is_dir()]
        assert len(sub_dirs) == 1
        assert re.match('[0-9]{4}-[0-9]{2}-[0-9]{2}_[0-9]{2}-[0-9]{2}-[0-9]{2}'
            , sub_dirs[0].name)

    @pytest.mark.unit
    def test_log_dir_overrides(self, monkeypatch, tmp_path):
        """Check a variety of trainer options with exp_manager"""
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {'explicit_log_dir': str(
            tmp_path / 'test_log_dir_overrides')})
        assert log_dir.resolve() == (tmp_path / 'test_log_dir_overrides'
            ).resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / 'test_log_dir_overrides').exists()
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {'exp_dir': str(tmp_path /
            'test_no_name'), 'version': 957})
        assert log_dir.resolve() == (tmp_path / 'test_no_name' / 'default' /
            '957').resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / 'test_no_name' / 'default' / '957').exists()
        monkeypatch.delenv(NEMO_ENV_VARNAME_VERSION)
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {'exp_dir': str(tmp_path /
            'test_no_name'), 'use_datetime_version': False})
        assert log_dir.resolve() == (tmp_path / 'test_no_name' / 'default' /
            'version_0').resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / 'test_no_name' / 'default' / 'version_0'
            ).exists()
        monkeypatch.delenv(NEMO_ENV_VARNAME_VERSION)
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {'exp_dir': str(tmp_path /
            'test_no_name'), 'use_datetime_version': False})
        assert log_dir.resolve() == (tmp_path / 'test_no_name' / 'default' /
            'version_1').resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / 'test_no_name' / 'default' / 'version_1'
            ).exists()

    @pytest.mark.unit
    def test_resume(self, tmp_path):
        """ Tests the resume capabilities of exp_manager"""
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        with pytest.raises(NotFoundError):
            exp_manager(test_trainer, {'exp_dir': str(tmp_path /
                'test_resume'), 'resume_if_exists': True,
                'explicit_log_dir': 'Does_not_exist'})
        with pytest.raises(NotFoundError):
            exp_manager(test_trainer, {'resume_if_exists': True, 'exp_dir':
                str(tmp_path / 'test_resume')})
        exp_manager(test_trainer, {'resume_if_exists': True, 'exp_dir': str
            (tmp_path / 'test_resume_2'), 'resume_ignore_no_checkpoint': True})
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' / 'checkpoints'
            ).mkdir(parents=True)
        with pytest.raises(NotFoundError):
            exp_manager(test_trainer, {'resume_if_exists': True,
                'explicit_log_dir': str(tmp_path / 'test_resume' /
                'default' / 'version_0')})
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel--end.ckpt').touch()
        with pytest.raises(ValueError):
            exp_manager(test_trainer, {'resume_if_exists': True,
                'explicit_log_dir': str(tmp_path / 'test_resume' /
                'default' / 'version_0')})
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel--end.ckpt').unlink()
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel--last.ckpt').touch()
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel2--last.ckpt').touch()
        with pytest.raises(ValueError):
            exp_manager(test_trainer, {'resume_if_exists': True,
                'explicit_log_dir': str(tmp_path / 'test_resume' /
                'default' / 'version_0')})
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel2--last.ckpt').unlink()
        log_dir = exp_manager(test_trainer, {'resume_if_exists': True,
            'explicit_log_dir': str(tmp_path / 'test_resume' / 'default' /
            'version_0')})
        checkpoint = Path(tmp_path / 'test_resume' / 'default' /
            'version_0' / 'checkpoints' / 'mymodel--last.ckpt')
        assert Path(test_trainer.resume_from_checkpoint).resolve(
            ) == checkpoint.resolve()
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        exp_manager(test_trainer, {'resume_if_exists': True,
            'explicit_log_dir': str(log_dir)})
        checkpoint = Path(tmp_path / 'test_resume' / 'default' /
            'version_0' / 'checkpoints' / 'mymodel--last.ckpt')
        assert Path(test_trainer.resume_from_checkpoint).resolve(
            ) == checkpoint.resolve()
        prev_run_dir = Path(tmp_path / 'test_resume' / 'default' /
            'version_0' / 'run_0')
        assert prev_run_dir.exists()
        prev_log = Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'run_0' / 'lightning_logs.txt')
        assert prev_log.exists()

import re
import shutil
from pathlib import Path
import pytest
import pytorch_lightning as pl
from omegaconf.errors import OmegaConfBaseException
from nemo.constants import NEMO_ENV_VARNAME_VERSION
from nemo.utils.exp_manager import CheckpointMisconfigurationError, LoggerMisconfigurationError, NotFoundError, exp_manager


@pytest.fixture
def cleanup_local_folder():
    assert not Path('./lightning_logs').exists()
    assert not Path('./NeMo_experiments').exists()
    assert not Path('./nemo_experiments').exists()
    yield
    if Path('./lightning_logs').exists():
        shutil.rmtree('./lightning_logs')
    if Path('./NeMo_experiments').exists():
        shutil.rmtree('./NeMo_experiments')
    if Path('./nemo_experiments').exists():
        shutil.rmtree('./nemo_experiments')


class TestExpManager:

    @pytest.mark.unit
    def test_omegaconf(self):
        """Ensure omegaconf raises an error when an unexcepted argument is passed"""
        with pytest.raises(OmegaConfBaseException):
            exp_manager(None, {'unused': 1})

    @pytest.mark.unit
    def test_trainer_loggers(self, cleanup_local_folder, tmp_path):
        """ Test that a trainer with logger errors out with a number of arguments. Test that it works with
        create_tensorboard_logger set to False
        """
        test_trainer = pl.Trainer()
        with pytest.raises(LoggerMisconfigurationError):
            exp_manager(test_trainer, {'exp_dir': str(tmp_path)})
        with pytest.raises(LoggerMisconfigurationError):
            exp_manager(test_trainer, {'explicit_log_dir': str(tmp_path)})
        with pytest.raises(LoggerMisconfigurationError):
            exp_manager(test_trainer, {'resume_if_exists': True})
        log_dir = exp_manager(test_trainer, {'create_tensorboard_logger': 
            False, 'create_checkpoint_callback': False})
        assert log_dir.resolve() == Path('./lightning_logs/version_0').resolve(
            )
        assert Path('./lightning_logs').exists()
        assert Path('./lightning_logs/version_0').exists()
        test_trainer = pl.Trainer(logger=False)
        log_dir = exp_manager(test_trainer, {'create_tensorboard_logger': 
            True, 'create_checkpoint_callback': False, 'exp_dir': str(
            tmp_path)})
        assert isinstance(test_trainer.logger, pl.loggers.TensorBoardLogger)
        test_trainer = pl.Trainer(logger=False)
        with pytest.raises(ValueError):
            log_dir = exp_manager(test_trainer, {
                'create_tensorboard_logger': False,
                'create_checkpoint_callback': False, 'exp_dir': str(
                tmp_path), 'create_wandb_logger': True})
        log_dir = exp_manager(test_trainer, {'create_tensorboard_logger': 
            False, 'create_checkpoint_callback': False, 'exp_dir': str(
            tmp_path), 'create_wandb_logger': True, 'wandb_logger_kwargs':
            {'name': '', 'project': ''}})
        assert isinstance(test_trainer.logger, pl.loggers.WandbLogger)

    @pytest.mark.unit
    def test_checkpoint_configurations(self, cleanup_local_folder):
        """ Test that trainer creating modelcheckpoint and asking exp_manager to do it too results in errors, but
        is error free if only one is asked to do so.
        """
        disable_tb_logger = {'create_tensorboard_logger': False}
        test_trainer = pl.Trainer()
        with pytest.raises(CheckpointMisconfigurationError):
            exp_manager(test_trainer, disable_tb_logger)
        exp_manager(test_trainer, {'create_checkpoint_callback': False,
            'create_tensorboard_logger': False})
        test_trainer_2 = pl.Trainer(checkpoint_callback=False)
        exp_manager(test_trainer_2, disable_tb_logger)

    @pytest.mark.unit
    def test_default_log_dir(self, cleanup_local_folder):
        """Check the default of ./nemo_experiments/default/datetime works as intended"""
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {'create_tensorboard_logger': 
            False, 'create_checkpoint_callback': False})
        assert (log_dir / '..').resolve() == Path('./nemo_experiments/default/'
            ).resolve()
        assert Path('./nemo_experiments').exists()
        assert Path('./nemo_experiments/default/').exists()
        sub_dirs = [x for x in Path('./nemo_experiments/default/').iterdir(
            ) if x.is_dir()]
        assert len(sub_dirs) == 1
        assert re.match('[0-9]{4}-[0-9]{2}-[0-9]{2}_[0-9]{2}-[0-9]{2}-[0-9]{2}'
            , sub_dirs[0].name)

    @pytest.mark.unit
    def test_log_dir_overrides(self, monkeypatch, tmp_path):
        """Check a variety of trainer options with exp_manager"""
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {'explicit_log_dir': str(
            tmp_path / 'test_log_dir_overrides')})
        assert log_dir.resolve() == (tmp_path / 'test_log_dir_overrides'
            ).resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / 'test_log_dir_overrides').exists()
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {'exp_dir': str(tmp_path /
            'test_no_name'), 'version': 957})
        assert log_dir.resolve() == (tmp_path / 'test_no_name' / 'default' /
            '957').resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / 'test_no_name' / 'default' / '957').exists()
        monkeypatch.delenv(NEMO_ENV_VARNAME_VERSION)
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {'exp_dir': str(tmp_path /
            'test_no_name'), 'use_datetime_version': False})
        assert log_dir.resolve() == (tmp_path / 'test_no_name' / 'default' /
            'version_0').resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / 'test_no_name' / 'default' / 'version_0'
            ).exists()
        monkeypatch.delenv(NEMO_ENV_VARNAME_VERSION)
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {'exp_dir': str(tmp_path /
            'test_no_name'), 'use_datetime_version': False})
        assert log_dir.resolve() == (tmp_path / 'test_no_name' / 'default' /
            'version_1').resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / 'test_no_name' / 'default' / 'version_1'
            ).exists()

    @pytest.mark.unit
    def test_resume(self, tmp_path):
        """ Tests the resume capabilities of exp_manager"""
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        with pytest.raises(NotFoundError):
            exp_manager(test_trainer, {'exp_dir': str(tmp_path /
                'test_resume'), 'resume_if_exists': True,
                'explicit_log_dir': 'Does_not_exist'})
        with pytest.raises(NotFoundError):
            exp_manager(test_trainer, {'resume_if_exists': True, 'exp_dir':
                str(tmp_path / 'test_resume')})
        exp_manager(test_trainer, {'resume_if_exists': True, 'exp_dir': str
            (tmp_path / 'test_resume_2'), 'resume_ignore_no_checkpoint': True})
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' / 'checkpoints'
            ).mkdir(parents=True)
        with pytest.raises(NotFoundError):
            exp_manager(test_trainer, {'resume_if_exists': True,
                'explicit_log_dir': str(tmp_path / 'test_resume' /
                'default' / 'version_0')})
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel--end.ckpt').touch()
        with pytest.raises(ValueError):
            exp_manager(test_trainer, {'resume_if_exists': True,
                'explicit_log_dir': str(tmp_path / 'test_resume' /
                'default' / 'version_0')})
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel--end.ckpt').unlink()
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel--last.ckpt').touch()
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel2--last.ckpt').touch()
        with pytest.raises(ValueError):
            exp_manager(test_trainer, {'resume_if_exists': True,
                'explicit_log_dir': str(tmp_path / 'test_resume' /
                'default' / 'version_0')})
        Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'checkpoints' / 'mymodel2--last.ckpt').unlink()
        log_dir = exp_manager(test_trainer, {'resume_if_exists': True,
            'explicit_log_dir': str(tmp_path / 'test_resume' / 'default' /
            'version_0')})
        checkpoint = Path(tmp_path / 'test_resume' / 'default' /
            'version_0' / 'checkpoints' / 'mymodel--last.ckpt')
        assert Path(test_trainer.resume_from_checkpoint).resolve(
            ) == checkpoint.resolve()
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        exp_manager(test_trainer, {'resume_if_exists': True,
            'explicit_log_dir': str(log_dir)})
        checkpoint = Path(tmp_path / 'test_resume' / 'default' /
            'version_0' / 'checkpoints' / 'mymodel--last.ckpt')
        assert Path(test_trainer.resume_from_checkpoint).resolve(
            ) == checkpoint.resolve()
        prev_run_dir = Path(tmp_path / 'test_resume' / 'default' /
            'version_0' / 'run_0')
        assert prev_run_dir.exists()
        prev_log = Path(tmp_path / 'test_resume' / 'default' / 'version_0' /
            'run_0' / 'lightning_logs.txt')
        assert prev_log.exists()

# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re
import shutil
from pathlib import Path

import pytest
import pytorch_lightning as pl
from omegaconf.errors import OmegaConfBaseException

from nemo.constants import NEMO_ENV_VARNAME_VERSION
from nemo.utils.exp_manager import (
    CheckpointMisconfigurationError,
    LoggerMisconfigurationError,
    NotFoundError,
    exp_manager,
)


@pytest.fixture
def cleanup_local_folder():
    # Asserts in fixture are not recommended, but I'd rather stop users from deleting expensive training runs
    assert not Path("./lightning_logs").exists()
    assert not Path("./NeMo_experiments").exists()
    assert not Path("./nemo_experiments").exists()

    yield

    if Path("./lightning_logs").exists():
        shutil.rmtree('./lightning_logs')
    if Path("./NeMo_experiments").exists():
        shutil.rmtree('./NeMo_experiments')
    if Path("./nemo_experiments").exists():
        shutil.rmtree('./nemo_experiments')


class TestExpManager:
    @pytest.mark.unit
    def test_omegaconf(self):
        """Ensure omegaconf raises an error when an unexcepted argument is passed"""
        with pytest.raises(OmegaConfBaseException):
            exp_manager(None, {"unused": 1})

    @pytest.mark.unit
    def test_trainer_loggers(self, cleanup_local_folder, tmp_path):
        """ Test that a trainer with logger errors out with a number of arguments. Test that it works with
        create_tensorboard_logger set to False
        """
        test_trainer = pl.Trainer()  # Should create logger and modelcheckpoint

        with pytest.raises(LoggerMisconfigurationError):  # Fails because exp_manager defaults to trainer
            exp_manager(test_trainer, {"exp_dir": str(tmp_path)})
        with pytest.raises(LoggerMisconfigurationError):  # Fails because exp_manager defaults to trainer
            exp_manager(test_trainer, {"explicit_log_dir": str(tmp_path)})
        with pytest.raises(LoggerMisconfigurationError):  # Fails because exp_manager defaults to trainer
            exp_manager(test_trainer, {"resume_if_exists": True})

        # Check that exp_manager uses trainer.logger, it's exp_dir, name, and version
        log_dir = exp_manager(test_trainer, {"create_tensorboard_logger": False, "create_checkpoint_callback": False})
        assert log_dir.resolve() == Path("./lightning_logs/version_0").resolve()
        assert Path("./lightning_logs").exists()
        assert Path("./lightning_logs/version_0").exists()

        # Check that a trainer without a logger gets a logger attached to it
        test_trainer = pl.Trainer(logger=False)
        log_dir = exp_manager(
            test_trainer,
            {"create_tensorboard_logger": True, "create_checkpoint_callback": False, "exp_dir": str(tmp_path)},
        )
        assert isinstance(test_trainer.logger, pl.loggers.TensorBoardLogger)

        test_trainer = pl.Trainer(logger=False)
        # Check that a create_wandb_logger=True errors out unless wandb_logger_kwargs is passed.
        with pytest.raises(ValueError):
            log_dir = exp_manager(
                test_trainer,
                {
                    "create_tensorboard_logger": False,
                    "create_checkpoint_callback": False,
                    "exp_dir": str(tmp_path),
                    "create_wandb_logger": True,
                },
            )
        # Check that a WandbLogger is attached to logger if create_wandb_logger=True and wandb_logger_kwargs has name
        # and project
        log_dir = exp_manager(
            test_trainer,
            {
                "create_tensorboard_logger": False,
                "create_checkpoint_callback": False,
                "exp_dir": str(tmp_path),
                "create_wandb_logger": True,
                "wandb_logger_kwargs": {"name": "", "project": ""},
            },
        )
        assert isinstance(test_trainer.logger, pl.loggers.WandbLogger)

    @pytest.mark.unit
    def test_checkpoint_configurations(self, cleanup_local_folder):
        """ Test that trainer creating modelcheckpoint and asking exp_manager to do it too results in errors, but
        is error free if only one is asked to do so.
        """
        disable_tb_logger = {"create_tensorboard_logger": False}
        test_trainer = pl.Trainer()  # Should create logger and modelcheckpoint
        with pytest.raises(CheckpointMisconfigurationError):  # Fails because both try to create modelcheckpoint
            exp_manager(test_trainer, disable_tb_logger)

        # Should succeed without error
        exp_manager(test_trainer, {"create_checkpoint_callback": False, "create_tensorboard_logger": False})

        test_trainer_2 = pl.Trainer(checkpoint_callback=False)
        exp_manager(test_trainer_2, disable_tb_logger)  # Should succeed without error

    @pytest.mark.unit
    def test_default_log_dir(self, cleanup_local_folder):
        """Check the default of ./nemo_experiments/default/datetime works as intended"""
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)

        log_dir = exp_manager(test_trainer, {"create_tensorboard_logger": False, "create_checkpoint_callback": False})
        assert (log_dir / "..").resolve() == Path("./nemo_experiments/default/").resolve()
        assert Path("./nemo_experiments").exists()
        assert Path("./nemo_experiments/default/").exists()
        sub_dirs = [x for x in Path("./nemo_experiments/default/").iterdir() if x.is_dir()]
        assert len(sub_dirs) == 1
        assert re.match(r"[0-9]{4}-[0-9]{2}-[0-9]{2}_[0-9]{2}-[0-9]{2}-[0-9]{2}", sub_dirs[0].name)

    @pytest.mark.unit
    def test_log_dir_overrides(self, monkeypatch, tmp_path):
        """Check a variety of trainer options with exp_manager"""
        # Checks that explicit_log_dir ignores exp_dir, name, and version
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {"explicit_log_dir": str(tmp_path / "test_log_dir_overrides")})
        assert log_dir.resolve() == (tmp_path / "test_log_dir_overrides").resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / "test_log_dir_overrides").exists()

        # Checks that exp_manager uses exp_dir, default name, and explicit version
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {"exp_dir": str(tmp_path / "test_no_name"), "version": 957})
        assert log_dir.resolve() == (tmp_path / "test_no_name" / "default" / "957").resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / "test_no_name" / "default" / "957").exists()

        monkeypatch.delenv(NEMO_ENV_VARNAME_VERSION)
        # Checks that use_datetime_version False toggle works
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {"exp_dir": str(tmp_path / "test_no_name"), "use_datetime_version": False})
        assert log_dir.resolve() == (tmp_path / "test_no_name" / "default" / "version_0").resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / "test_no_name" / "default" / "version_0").exists()

        monkeypatch.delenv(NEMO_ENV_VARNAME_VERSION)
        # Checks that use_datetime_version False toggle works and version increments
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        log_dir = exp_manager(test_trainer, {"exp_dir": str(tmp_path / "test_no_name"), "use_datetime_version": False})
        assert log_dir.resolve() == (tmp_path / "test_no_name" / "default" / "version_1").resolve()
        assert Path(tmp_path).exists()
        assert Path(tmp_path / "test_no_name" / "default" / "version_1").exists()

    @pytest.mark.unit
    def test_resume(self, tmp_path):
        """ Tests the resume capabilities of exp_manager"""
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)

        # Error because explicit_log_dir does not exist
        with pytest.raises(NotFoundError):
            exp_manager(
                test_trainer,
                {
                    "exp_dir": str(tmp_path / "test_resume"),
                    "resume_if_exists": True,
                    "explicit_log_dir": "Does_not_exist",
                },
            )

        # Error because checkpoints folder does not exist
        with pytest.raises(NotFoundError):
            exp_manager(test_trainer, {"resume_if_exists": True, "exp_dir": str(tmp_path / "test_resume")})

        # No error because we tell exp_manager to ignore notfounderror
        exp_manager(
            test_trainer,
            {
                "resume_if_exists": True,
                "exp_dir": str(tmp_path / "test_resume_2"),
                "resume_ignore_no_checkpoint": True,
            },
        )

        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints").mkdir(parents=True)
        # Error because checkpoints do not exist in folder
        with pytest.raises(NotFoundError):
            exp_manager(
                test_trainer,
                {
                    "resume_if_exists": True,
                    "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0"),
                },
            )

        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--end.ckpt").touch()
        # Error because *end.ckpt is in folder indicating that training has already finished
        with pytest.raises(ValueError):
            exp_manager(
                test_trainer,
                {
                    "resume_if_exists": True,
                    "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0"),
                },
            )

        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--end.ckpt").unlink()
        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--last.ckpt").touch()
        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel2--last.ckpt").touch()
        # Error because multiple *last.ckpt is in folder. If more than one, don't know which to restore
        with pytest.raises(ValueError):
            exp_manager(
                test_trainer,
                {
                    "resume_if_exists": True,
                    "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0"),
                },
            )

        # Finally succeed
        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel2--last.ckpt").unlink()
        log_dir = exp_manager(
            test_trainer,
            {"resume_if_exists": True, "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0")},
        )
        checkpoint = Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--last.ckpt")
        assert Path(test_trainer.resume_from_checkpoint).resolve() == checkpoint.resolve()

        # Succeed again and make sure that run_0 exists and previous log files were moved
        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
        exp_manager(test_trainer, {"resume_if_exists": True, "explicit_log_dir": str(log_dir)})
        checkpoint = Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--last.ckpt")
        assert Path(test_trainer.resume_from_checkpoint).resolve() == checkpoint.resolve()
        prev_run_dir = Path(tmp_path / "test_resume" / "default" / "version_0" / "run_0")
        assert prev_run_dir.exists()
        prev_log = Path(tmp_path / "test_resume" / "default" / "version_0" / "run_0" / "lightning_logs.txt")
        assert prev_log.exists()
