From 571898bc0f8ce3f3d9f83a779a7c6a44f485e1f2 Mon Sep 17 00:00:00 2001
From: M1nd3r <petrsedlacek.km@seznam.cz>
Date: Mon, 14 Oct 2024 18:37:29 +0200
Subject: [PATCH] chore(trezorlib): adjust benchmark feature for THP [no
 changelog]

---
 python/src/trezorlib/benchmark.py     | 10 +++++-----
 python/src/trezorlib/cli/benchmark.py | 24 +++++++++++-------------
 python/src/trezorlib/cli/trezorctl.py |  2 +-
 3 files changed, 17 insertions(+), 19 deletions(-)

diff --git a/python/src/trezorlib/benchmark.py b/python/src/trezorlib/benchmark.py
index f96ef7970e..b961dda426 100644
--- a/python/src/trezorlib/benchmark.py
+++ b/python/src/trezorlib/benchmark.py
@@ -20,17 +20,17 @@ from . import messages
 from .tools import expect
 
 if TYPE_CHECKING:
-    from .client import TrezorClient
     from .protobuf import MessageType
+    from .transport.session import Session
 
 
 @expect(messages.BenchmarkNames)
 def list_names(
-    client: "TrezorClient",
+    session: "Session",
 ) -> "MessageType":
-    return client.call(messages.BenchmarkListNames())
+    return session.call(messages.BenchmarkListNames())
 
 
 @expect(messages.BenchmarkResult)
-def run(client: "TrezorClient", name: str) -> "MessageType":
-    return client.call(messages.BenchmarkRun(name=name))
+def run(session: "Session", name: str) -> "MessageType":
+    return session.call(messages.BenchmarkRun(name=name))
diff --git a/python/src/trezorlib/cli/benchmark.py b/python/src/trezorlib/cli/benchmark.py
index e445089815..a7ebab12f7 100644
--- a/python/src/trezorlib/cli/benchmark.py
+++ b/python/src/trezorlib/cli/benchmark.py
@@ -20,17 +20,15 @@ from typing import TYPE_CHECKING, List, Optional
 import click
 
 from .. import benchmark
-from . import with_client
+from . import with_session
 
 if TYPE_CHECKING:
 
-    from ..client import TrezorClient
+    from ..transport.session import Session
 
 
-def list_names_patern(
-    client: "TrezorClient", pattern: Optional[str] = None
-) -> List[str]:
-    names = list(benchmark.list_names(client).names)
+def list_names_patern(session: "Session", pattern: Optional[str] = None) -> List[str]:
+    names = list(benchmark.list_names(session).names)
     if pattern is None:
         return names
     return [name for name in names if fnmatch(name, pattern)]
@@ -43,10 +41,10 @@ def cli() -> None:
 
 @cli.command()
 @click.argument("pattern", required=False)
-@with_client
-def list_names(client: "TrezorClient", pattern: Optional[str] = None) -> None:
+@with_session
+def list_names(session: "Session", pattern: Optional[str] = None) -> None:
     """List names of all supported benchmarks"""
-    names = list_names_patern(client, pattern)
+    names = list_names_patern(session, pattern)
     if len(names) == 0:
         click.echo("No benchmark satisfies the pattern.")
     else:
@@ -56,13 +54,13 @@ def list_names(client: "TrezorClient", pattern: Optional[str] = None) -> None:
 
 @cli.command()
 @click.argument("pattern", required=False)
-@with_client
-def run(client: "TrezorClient", pattern: Optional[str]) -> None:
+@with_session
+def run(session: "Session", pattern: Optional[str]) -> None:
     """Run benchmark"""
-    names = list_names_patern(client, pattern)
+    names = list_names_patern(session, pattern)
     if len(names) == 0:
         click.echo("No benchmark satisfies the pattern.")
     else:
         for name in names:
-            result = benchmark.run(client, name)
+            result = benchmark.run(session, name)
             click.echo(f"{name}: {result.value} {result.unit}")
diff --git a/python/src/trezorlib/cli/trezorctl.py b/python/src/trezorlib/cli/trezorctl.py
index f0d2b07c9f..640adc3273 100755
--- a/python/src/trezorlib/cli/trezorctl.py
+++ b/python/src/trezorlib/cli/trezorctl.py
@@ -32,8 +32,8 @@ from ..transport.session import Session
 from ..transport.udp import UdpTransport
 from . import (
     AliasedGroup,
-    benchmark,
     NewTrezorConnection,
+    benchmark,
     binance,
     btc,
     cardano,