From 71c078f24c5de7efe9ca30a73ae829abdc44cccd Mon Sep 17 00:00:00 2001
From: obrusvit <vit.obrusnik@gmail.com>
Date: Tue, 11 Mar 2025 16:04:15 +0100
Subject: [PATCH] fix(tests): handle layout-specific translations

---
 python/src/trezorlib/debuglink.py |  8 ++++++--
 tests/translations.py             | 22 ++++++++++++++++++++--
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/python/src/trezorlib/debuglink.py b/python/src/trezorlib/debuglink.py
index 77c5be7bf2..09e055f6df 100644
--- a/python/src/trezorlib/debuglink.py
+++ b/python/src/trezorlib/debuglink.py
@@ -93,10 +93,14 @@ class LayoutType(Enum):
             return cls.Eckhart
         if model in (models.T1B1,):
             return cls.T1
-        if model in (models.T3W1,):
-            return cls.Eckhart
         raise ValueError(f"Unknown model: {model}")
 
+    def __str__(self) -> str:
+        return self.name
+
+    def __repr__(self) -> str:
+        return f"LayoutType.{self.name}"
+
 
 class UnstructuredJSONReader:
     """Contains data-parsing helpers for JSON data that have unknown structure."""
diff --git a/tests/translations.py b/tests/translations.py
index 84d8a95f27..33e8589e75 100644
--- a/tests/translations.py
+++ b/tests/translations.py
@@ -8,6 +8,7 @@ from pathlib import Path
 
 from trezorlib import cosi, device, models
 from trezorlib._internal import translations
+from trezorlib.debuglink import LayoutType
 from trezorlib.debuglink import TrezorClientDebugLink as Client
 
 from . import common
@@ -72,6 +73,11 @@ def set_language(client: Client, lang: str):
     with client:
         device.change_language(client, language_data)  # type: ignore
     _CURRENT_TRANSLATION.TR = TRANSLATIONS[lang]
+    _CURRENT_TRANSLATION.LAYOUT = client.layout_type
+
+
+def set_layout(client: Client):
+    _CURRENT_TRANSLATION.LAYOUT = client.layout_type
 
 
 def get_lang_json(lang: str) -> translations.JsonDef:
@@ -91,13 +97,24 @@ class Translation:
         self.lang_json = get_lang_json(lang)
 
     @property
-    def translations(self) -> dict[str, str]:
+    def translations(self) -> dict[str, str | dict[str, str]]:
         return self.lang_json["translations"]
 
     def _translate_raw(self, key: str, _stacklevel: int = 0) -> str:
         tr = self.translations.get(key)
         if tr is not None:
-            return tr
+            # Handle layout-specific translations
+            if isinstance(tr, dict) and hasattr(_CURRENT_TRANSLATION, "LAYOUT"):
+                # Try to get translation for current layout
+                layout_name = _CURRENT_TRANSLATION.LAYOUT.name
+                if layout_name in tr:
+                    return tr[layout_name]
+                # Fall back to any available translation if no match for current layout
+                return next(iter(tr.values()))
+            elif isinstance(tr, str):
+                return tr
+            else:
+                raise ValueError(f"Invalid translation value for key '{key}'")
         if self.lang != "en":
             # check if the key exists in English first
             retval = TRANSLATIONS["en"]._translate_raw(key)
@@ -122,6 +139,7 @@ class Translation:
 
 TRANSLATIONS = {lang: Translation(lang) for lang in LANGUAGES}
 _CURRENT_TRANSLATION.TR = TRANSLATIONS["en"]
+_CURRENT_TRANSLATION.LAYOUT = LayoutType.Bolt
 
 
 def translate(key: str, _stacklevel: int = 0) -> str: