From c98b0c4f082497adeb572daad70374be69ac8152 Mon Sep 17 00:00:00 2001
From: Ethan Harris <ethanwharris@gmail.com>
Date: Tue, 20 Dec 2022 20:46:38 +0000
Subject: [PATCH] [App] Fix e2e tests (#16146)

---
 src/lightning_app/testing/testing.py | 40 +++++++++++-----------------
 1 file changed, 15 insertions(+), 25 deletions(-)

diff --git a/src/lightning_app/testing/testing.py b/src/lightning_app/testing/testing.py
index 0c73f79b59..5602af1e52 100644
--- a/src/lightning_app/testing/testing.py
+++ b/src/lightning_app/testing/testing.py
@@ -386,34 +386,24 @@ def run_app_in_cloud(
             process = Process(target=_print_logs, kwargs={"app_id": app_id})
             process.start()
 
-        if not app.spec.is_headless:
-            while True:
-                try:
-                    with admin_page.context.expect_page() as page_catcher:
-                        admin_page.locator('[data-cy="open"]').click()
-                    view_page = page_catcher.value
-                    view_page.wait_for_load_state(timeout=0)
-                    break
-                except (playwright._impl._api_types.Error, playwright._impl._api_types.TimeoutError):
-                    pass
-        else:
-            view_page = None
+        # Wait until the app is running
+        while True:
+            sleep(1)
 
-            # Wait until the app is running
-            while True:
-                sleep(1)
+            lit_apps = [
+                app
+                for app in client.lightningapp_instance_service_list_lightningapp_instances(
+                    project_id=project.project_id
+                ).lightningapps
+                if app.name == name
+            ]
+            app = lit_apps[0]
 
-                lit_apps = [
-                    app
-                    for app in client.lightningapp_instance_service_list_lightningapp_instances(
-                        project_id=project.project_id
-                    ).lightningapps
-                    if app.name == name
-                ]
-                app = lit_apps[0]
+            if app.status.phase == V1LightningappInstanceState.RUNNING:
+                break
 
-                if app.status.phase == V1LightningappInstanceState.RUNNING:
-                    break
+        view_page = context.new_page()
+        view_page.goto(f"{app.status.url}/view")
 
         # TODO: is re-creating this redundant?
         lit_apps = [