experimental: split exit method and logic for screenshots capture

Author: anastasiyaig

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
@@ -118,10 +129,10 @@ def startaut(self):
 
     @allure.step('Close application')
     def stop(self):
+        time.sleep(1)  # Some time before stopping the app to capture screenshot.
         LOG.info('Stopping AUT: %s', self.path)
         self.detach_context()
         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':