From f22f4d920e2b2cca557852ccbc06a4da5dc28de4 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Mon, 21 Oct 2019 16:49:34 +0200
Subject: [PATCH] tests: ensure Shamir persistence test runs only for core

---
 tests/device_handler.py                       |  3 +-
 tests/upgrade_tests/__init__.py               | 37 +++++++++++++++++++
 tests/upgrade_tests/test_firmware_upgrades.py | 10 ++---
 .../upgrade_tests/test_shamir_persistence.py  |  2 +
 4 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/tests/device_handler.py b/tests/device_handler.py
index 4080c498a3..11969841d2 100644
--- a/tests/device_handler.py
+++ b/tests/device_handler.py
@@ -72,5 +72,6 @@ class BackgroundDeviceHandler:
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        if not self.check_finalize():
+        finalized_ok = self.check_finalize()
+        if exc_type is None and not finalized_ok:
             raise RuntimeError("Exit while task is unfinished")
diff --git a/tests/upgrade_tests/__init__.py b/tests/upgrade_tests/__init__.py
index e69de29bb2..efeb28409d 100644
--- a/tests/upgrade_tests/__init__.py
+++ b/tests/upgrade_tests/__init__.py
@@ -0,0 +1,37 @@
+import os
+
+import pytest
+
+from ..emulators import EmulatorWrapper
+
+SELECTED_GENS = [
+    gen.strip() for gen in os.environ.get("TREZOR_UPGRADE_TEST", "").split(",") if gen
+]
+
+if SELECTED_GENS:
+    # if any gens were selected via the environment variable, force enable all selected
+    LEGACY_ENABLED = "legacy" in SELECTED_GENS
+    CORE_ENABLED = "core" in SELECTED_GENS
+
+else:
+    # if no selection was provided, select those for which we have emulators
+    try:
+        EmulatorWrapper("legacy")
+        LEGACY_ENABLED = True
+    except Exception:
+        LEGACY_ENABLED = False
+
+    try:
+        EmulatorWrapper("core")
+        CORE_ENABLED = True
+    except Exception:
+        CORE_ENABLED = False
+
+
+legacy_only = pytest.mark.skipif(
+    not LEGACY_ENABLED, reason="This test requires legacy emulator"
+)
+
+core_only = pytest.mark.skipif(
+    not CORE_ENABLED, reason="This test requires core emulator"
+)
diff --git a/tests/upgrade_tests/test_firmware_upgrades.py b/tests/upgrade_tests/test_firmware_upgrades.py
index bea99e2900..39f3d078bf 100644
--- a/tests/upgrade_tests/test_firmware_upgrades.py
+++ b/tests/upgrade_tests/test_firmware_upgrades.py
@@ -14,14 +14,13 @@
 # You should have received a copy of the License along with this library.
 # If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.
 
-import os
-
 import pytest
 
 from trezorlib import MINIMUM_FIRMWARE_VERSION, btc, debuglink, device
 from trezorlib.tools import H_
 
 from ..emulators import ALL_TAGS, EmulatorWrapper
+from . import SELECTED_GENS
 
 MINIMUM_FIRMWARE_VERSION["1"] = (1, 0, 0)
 MINIMUM_FIRMWARE_VERSION["T"] = (2, 0, 0)
@@ -41,11 +40,8 @@ def for_all(*args, minimum_version=(1, 0, 0)):
     if not args:
         args = ("core", "legacy")
 
-    specified_gens = os.environ.get("TREZOR_UPGRADE_TEST")
-    if specified_gens is not None:
-        enabled_gens = specified_gens.split(",")
-    else:
-        enabled_gens = args
+    # If any gens were selected, use them. If none, select all.
+    enabled_gens = SELECTED_GENS or args
 
     all_params = []
     for gen in args:
diff --git a/tests/upgrade_tests/test_shamir_persistence.py b/tests/upgrade_tests/test_shamir_persistence.py
index 923a6c2cf4..8cfea04e90 100644
--- a/tests/upgrade_tests/test_shamir_persistence.py
+++ b/tests/upgrade_tests/test_shamir_persistence.py
@@ -19,6 +19,7 @@ from trezorlib import device
 from .. import buttons
 from ..device_handler import BackgroundDeviceHandler
 from ..emulators import EmulatorWrapper
+from . import core_only
 
 
 def enter_word(debug, word):
@@ -28,6 +29,7 @@ def enter_word(debug, word):
     return debug.click(buttons.CONFIRM_WORD, wait=True)
 
 
+@core_only
 def test_persistence():
     with EmulatorWrapper("core") as emu, BackgroundDeviceHandler(
         emu.client