From 06fe820b064bcee86178ae462786f397c02ad8bb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Carlos=20Mochol=C3=AD?= <carlossmocholi@gmail.com>
Date: Thu, 15 Dec 2022 21:16:13 +0100
Subject: [PATCH] Further fixes

---
 .../public/test_multi_node.py                 | 48 +++++++------------
 1 file changed, 17 insertions(+), 31 deletions(-)

diff --git a/tests/tests_examples_app/public/test_multi_node.py b/tests/tests_examples_app/public/test_multi_node.py
index e664fb40b21c6f..2cb8d3b7690fdd 100644
--- a/tests/tests_examples_app/public/test_multi_node.py
+++ b/tests/tests_examples_app/public/test_multi_node.py
@@ -1,12 +1,11 @@
 import os
-import sys
 from unittest import mock
 
 import pytest
-from lightning_utilities.core.imports import module_available
 from tests_examples_app.public import _PATH_EXAMPLES
 
-from lightning_app.testing.testing import application_testing, LightningTestApp
+from lightning.app.testing.testing import application_testing, LightningTestApp
+from lightning_app.testing.helpers import _RunIf
 
 
 class LightningTestMultiNodeApp(LightningTestApp):
@@ -18,39 +17,26 @@ def on_before_run_once(self):
         return res
 
 
-@pytest.mark.skip(reason="flaky")
-@mock.patch("lightning_app.components.multi_node.base.is_running_in_cloud", return_value=True)
-def test_multi_node_example(_, monkeypatch):
-    monkeypatch.chdir(os.path.join(_PATH_EXAMPLES, "app_multi_node"))
-    command_line = [
-        "app.py",
-        "--blocking",
-        "False",
-        "--open-ui",
-        "False",
-    ]
-    result = application_testing(LightningTestMultiNodeApp, command_line)
-    assert result.exit_code == 0
-
-
-class LightningTestMultiNodeWorksApp(LightningTestApp):
-    def on_before_run_once(self):
-        res = super().on_before_run_once()
-        if self.works and all(w.has_stopped for w in self.works):
-            assert len([w for w in self.works]) == 2
-            return True
-        return res
-
-
 @pytest.mark.parametrize(
     "app_name",
-    ["train_pytorch.py", "train_any.py", "train_lite.py", "train_pytorch_spawn.py", "train_lt.py"],
+    [
+        "train_lt_script.py",
+        "train_pytorch.py",
+        "train_any.py",
+        "train_lite.py",
+        "train_pytorch_spawn.py",
+        "train_lt.py",
+    ],
 )
-@pytest.mark.skipif(sys.platform == "win32", reason="flaky")
-@pytest.mark.skipif(not module_available("lightning"), reason="lightning not available")
+@_RunIf(skip_windows=True)  # flaky
 @mock.patch("lightning_app.components.multi_node.base.is_running_in_cloud", return_value=True)
 def test_multi_node_examples(_, app_name, monkeypatch):
+    # note: this test will fail locally:
+    # * if you installed `lightning_app`, then the examples need to be
+    # rewritten to use `lightning_app` imports (CI does this)
+    # * if you installed `lightning`, then the imports in this file and mocks
+    # need to be changed to use `lightning`.
     monkeypatch.chdir(os.path.join(_PATH_EXAMPLES, "app_multi_node"))
     command_line = [app_name, "--blocking", "False", "--open-ui", "False", "--setup"]
-    result = application_testing(LightningTestMultiNodeWorksApp, command_line)
+    result = application_testing(LightningTestMultiNodeApp, command_line)
     assert result.exit_code == 0