lightning/tests/tests_app/utilities/test_imports.py

import os
from unittest import mock

import pytest

from lightning_app.utilities.imports import requires


@mock.patch.dict(os.environ, {"LIGHTING_TESTING": "0"})
def test_requires():
    @requires("lightning_app")
    def fn():
        pass

    fn()

    @requires("shouldnotexist")
    def fn_raise():
        pass

    with pytest.raises(ModuleNotFoundError, match="Please run: pip install 'shouldnotexist'"):
        fn_raise()

    class ClassRaise:
        @requires("shouldnotexist")
        def __init__(self):
            pass

    with pytest.raises(ModuleNotFoundError, match="Please run: pip install 'shouldnotexist'"):
        ClassRaise()


@mock.patch.dict(os.environ, {"LIGHTING_TESTING": "0"})
def test_requires_multiple():
    @requires(["shouldnotexist1", "shouldnotexist2"])
    def fn_raise():
        pass

    with pytest.raises(ModuleNotFoundError, match="Please run: pip install 'shouldnotexist1' 'shouldnotexist2'"):
        fn_raise()