From 1c30adb5484fe9fbdc63fd82522d1c1f6115c7f2 Mon Sep 17 00:00:00 2001
From: M1nd3r <petrsedlacek.km@seznam.cz>
Date: Tue, 4 Feb 2025 15:13:35 +0100
Subject: [PATCH] chore(tests): adapt testing framework to session based [no
 changelog]

---
 tests/conftest.py | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 2cf5d4b394..fc79d9824e 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -22,6 +22,7 @@ from enum import IntEnum
 from pathlib import Path
 from time import sleep
 
+import cryptography
 import pytest
 import xdist
 from _pytest.python import IdMaker
@@ -303,11 +304,23 @@ def _client_unlocked(
         should_format = sd_marker.kwargs.get("formatted", True)
         _raw_client.debug.erase_sd_card(format=should_format)
 
-    if _raw_client.is_invalidated:
-        _raw_client = _raw_client.get_new_client()
-    session = _raw_client.get_seedless_session()
-    wipe_device(session)
-    sleep(1.5)  # Makes tests more stable (wait for wipe to finish)
+    while True:
+        try:
+            if _raw_client.is_invalidated:
+                _raw_client = _raw_client.get_new_client()
+            session = _raw_client.get_seedless_session()
+            wipe_device(session)
+            sleep(1.5)  # Makes tests more stable (wait for wipe to finish)
+            break
+        except cryptography.exceptions.InvalidTag:
+            # Get a new client
+            _raw_client = _get_raw_client(request)
+
+    _raw_client.protocol = None
+    _raw_client.__init__(
+        transport=_raw_client.transport,
+        auto_interact=_raw_client.debug.allow_interactions,
+    )
 
     if not _raw_client.features.bootloader_mode:
         _raw_client.refresh_features()