From 461f54af0eec2e48736c994459a3dcba6f908536 Mon Sep 17 00:00:00 2001
From: Anastasiya <anastasija.ig@gmail.com>
Date: Fri, 7 Feb 2025 17:50:37 +0300
Subject: [PATCH] experimental: split exit method and logic for screenshots
 capture

---
 test/e2e/driver/aut.py | 53 +++++++++++++++++++++++++-----------------
 1 file changed, 32 insertions(+), 21 deletions(-)

diff --git a/test/e2e/driver/aut.py b/test/e2e/driver/aut.py
index 65ebb35d88e..cddf99eda51 100644
--- a/test/e2e/driver/aut.py
+++ b/test/e2e/driver/aut.py
@@ -47,28 +47,39 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_value, traceback):
         if exc_type:
-            try:
-                self.attach()
-                driver.waitForObjectExists(statusDesktop_mainWindow).setVisible(True)
-                configs.testpath.TEST.mkdir(parents=True, exist_ok=True)
-                screenshot = configs.testpath.TEST / f'{self.aut_id}.png'
-
-                rect = driver.object.globalBounds(driver.waitForObject(statusDesktop_mainWindow))
-                img = ImageGrab.grab(
-                    bbox=(rect.x, rect.y, rect.x + rect.width, rect.y + rect.height),
-                    xdisplay=configs.system.DISPLAY if get_platform() == "Linux" else None)
-                view = cv2.cvtColor(np.array(img), cv2.COLOR_BGR2RGB)
-                cv2.imwrite(str(screenshot), view)
-
-                allure.attach(
-                    name=f'Screenshot on fail: {self.aut_id}',
-                    body=screenshot.read_bytes(),
-                    attachment_type=allure.attachment_type.PNG)
-            except Exception as err:
-                LOG.error(err)
-
+            self._capture_screenshot()
         self.stop()
 
+    def _capture_screenshot(self):
+        """Captures a screenshot of the main application window if it is still running."""
+        try:
+            if driver.isRunning(self.aut_id):  # Проверяем, работает ли приложение
+                main_window = driver.waitForObjectExists(statusDesktop_mainWindow, 500)
+                if main_window:
+                    rect = driver.object.globalBounds(main_window)
+                    screenshot = configs.testpath.TEST / f'{self.aut_id}.png'
+                    configs.testpath.TEST.mkdir(parents=True, exist_ok=True)
+
+                    img = ImageGrab.grab(
+                        bbox=(rect.x, rect.y, rect.x + rect.width, rect.y + rect.height),
+                        xdisplay=configs.system.DISPLAY if get_platform() == "Linux" else None
+                    )
+                    view = cv2.cvtColor(np.array(img), cv2.COLOR_BGR2RGB)
+                    cv2.imwrite(str(screenshot), view)
+
+                    allure.attach(
+                        name=f'Screenshot on fail: {self.aut_id}',
+                        body=screenshot.read_bytes(),
+                        attachment_type=allure.attachment_type.PNG
+                    )
+                    LOG.info(f"Screenshot saved at {screenshot}")
+                else:
+                    LOG.warning("Main window not found; skipping screenshot.")
+            else:
+                LOG.warning("Application is no longer running; skipping screenshot.")
+        except Exception as err:
+            LOG.error(f"Failed to capture screenshot: {err}")
+
     def detach_context(self):
         if self.ctx is None:
             return
@@ -120,8 +131,8 @@ def startaut(self):
     def stop(self):
         LOG.info('Stopping AUT: %s', self.path)
         self.detach_context()
+        time.sleep(2)  # some time before kiling process, to allow screenshot capturing
         local_system.kill_process(self.pid)
-        time.sleep(1) # FIXME: Implement waiting for process to actually exit.
 
     @allure.step("Start and attach AUT")
     def launch(self) -> 'AUT':