From 508b6a87e6df26c1ea2b6e9c2a7f43fb22aef5f7 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Thu, 19 Jan 2023 13:35:19 +0100
Subject: [PATCH] fix(tests): ensure the client is open when stopping screen
 recording

[no changelog]
---
 python/src/trezorlib/debuglink.py | 5 +++++
 tests/ui_tests/__init__.py        | 1 +
 2 files changed, 6 insertions(+)

diff --git a/python/src/trezorlib/debuglink.py b/python/src/trezorlib/debuglink.py
index a05c38809e..c9ef48c5b0 100644
--- a/python/src/trezorlib/debuglink.py
+++ b/python/src/trezorlib/debuglink.py
@@ -618,6 +618,11 @@ class TrezorClientDebugLink(TrezorClient):
             Callable[[protobuf.MessageType], protobuf.MessageType],
         ] = {}
 
+    def ensure_open(self) -> None:
+        """Only open session if there isn't already an open one."""
+        if self.session_counter == 0:
+            self.open()
+
     def open(self) -> None:
         super().open()
         if self.session_counter == 1:
diff --git a/tests/ui_tests/__init__.py b/tests/ui_tests/__init__.py
index eed6927291..fe2d1be73f 100644
--- a/tests/ui_tests/__init__.py
+++ b/tests/ui_tests/__init__.py
@@ -153,6 +153,7 @@ def screen_recording(
         client.debug.start_recording(str(screen_path))
         yield
     finally:
+        client.ensure_open()
         # Wait for response to Initialize, which gives the emulator time to catch up
         # and redraw the homescreen. Otherwise there's a race condition between that
         # and stopping recording.