From 0d798211455b59be55b15e60d953b07bbeaaee0f Mon Sep 17 00:00:00 2001
From: Joost van Doorn <joost.van.doorn@gmail.com>
Date: Sun, 20 Mar 2022 08:58:32 +0100
Subject: [PATCH] Test correctness of module flags

---
 tests/utilities/test_imports.py | 56 ++++++++++++++++++++++++++++++++-
 1 file changed, 55 insertions(+), 1 deletion(-)

diff --git a/tests/utilities/test_imports.py b/tests/utilities/test_imports.py
index 75bcb51ffb89f..8f1937d0512eb 100644
--- a/tests/utilities/test_imports.py
+++ b/tests/utilities/test_imports.py
@@ -13,7 +13,16 @@
 # limitations under the License.
 import operator
 
-from pytorch_lightning.utilities import _module_available
+from pytorch_lightning.utilities import (
+    _APEX_AVAILABLE,
+    _BAGUA_AVAILABLE,
+    _DEEPSPEED_AVAILABLE,
+    _FAIRSCALE_AVAILABLE,
+    _HOROVOD_AVAILABLE,
+    _module_available,
+    _OMEGACONF_AVAILABLE,
+    _POPTORCH_AVAILABLE,
+)
 from pytorch_lightning.utilities.imports import _compare_version
 
 
@@ -45,3 +54,48 @@ def test_compare_version(monkeypatch):
     assert not _compare_version("torch", operator.ge, "1.10.0.rc0")
     assert _compare_version("torch", operator.ge, "1.10.0", use_base_version=True)
     assert not _compare_version("torch", operator.ge, "1.10.0")
+
+
+def test_imports():
+    try:
+        import apex  # noqa
+
+        assert _APEX_AVAILABLE
+    except ModuleNotFoundError:
+        assert not _APEX_AVAILABLE
+    try:
+        import bagua  # noqa
+
+        assert _BAGUA_AVAILABLE
+    except ModuleNotFoundError:
+        assert not _BAGUA_AVAILABLE
+    try:
+        import deepspeed  # noqa
+
+        assert _DEEPSPEED_AVAILABLE
+    except ModuleNotFoundError:
+        assert not _DEEPSPEED_AVAILABLE
+    try:
+        import fairscale.nn  # noqa
+
+        assert _FAIRSCALE_AVAILABLE
+    except ModuleNotFoundError:
+        assert not _FAIRSCALE_AVAILABLE
+    try:
+        import horovod.torch  # noqa
+
+        assert _HOROVOD_AVAILABLE
+    except ModuleNotFoundError:
+        assert not _HOROVOD_AVAILABLE
+    try:
+        import omegaconf  # noqa
+
+        assert _OMEGACONF_AVAILABLE
+    except ModuleNotFoundError:
+        assert not _OMEGACONF_AVAILABLE
+    try:
+        import poptorch  # noqa
+
+        assert _POPTORCH_AVAILABLE
+    except ModuleNotFoundError:
+        assert not _POPTORCH_AVAILABLE