From f7c04bcef84bb71b55c7018764791fe138104917 Mon Sep 17 00:00:00 2001
From: Steven B <51370195+sdb9696@users.noreply.github.com>
Date: Wed, 24 Jan 2024 09:40:36 +0000
Subject: [PATCH] Add --batch-size hint to timeout errors in dump_devinfo
 (#696)

* Add --batch-size hint to timeout errors in dump_devinfo

* Add _echo_error function for displaying critical errors
---
 devtools/dump_devinfo.py | 57 ++++++++++++++++++++++++++--------------
 1 file changed, 37 insertions(+), 20 deletions(-)

diff --git a/devtools/dump_devinfo.py b/devtools/dump_devinfo.py
index 6a8240ef..e9ec56b7 100644
--- a/devtools/dump_devinfo.py
+++ b/devtools/dump_devinfo.py
@@ -20,7 +20,14 @@ from typing import Dict, List, Union
 import asyncclick as click
 
 from devtools.helpers.smartrequests import COMPONENT_REQUESTS, SmartRequest
-from kasa import AuthenticationException, Credentials, Discover, SmartDevice
+from kasa import (
+    AuthenticationException,
+    Credentials,
+    Discover,
+    SmartDevice,
+    SmartDeviceException,
+    TimeoutException,
+)
 from kasa.discover import DiscoveryResult
 from kasa.exceptions import SmartErrorCode
 from kasa.tapo.tapodevice import TapoDevice
@@ -227,11 +234,7 @@ async def get_legacy_fixture(device):
     try:
         final = await device.protocol.query(final_query)
     except Exception as ex:
-        click.echo(
-            click.style(
-                f"Unable to query all successes at once: {ex}", bold=True, fg="red"
-            )
-        )
+        _echo_error(f"Unable to query all successes at once: {ex}", bold=True, fg="red")
 
     if device._discovery_info and not device._discovery_info.get("system"):
         # Need to recreate a DiscoverResult here because we don't want the aliases
@@ -254,6 +257,16 @@ async def get_legacy_fixture(device):
     return save_filename, copy_folder, final
 
 
+def _echo_error(msg: str):
+    click.echo(
+        click.style(
+            msg,
+            bold=True,
+            fg="red",
+        )
+    )
+
+
 async def _make_requests_or_exit(
     device: SmartDevice,
     requests: List[SmartRequest],
@@ -277,17 +290,25 @@ async def _make_requests_or_exit(
                 final[method] = result
         return final
     except AuthenticationException as ex:
-        click.echo(
-            click.style(
-                f"Unable to query the device due to an authentication error: {ex}",
-                bold=True,
-                fg="red",
-            )
+        _echo_error(
+            f"Unable to query the device due to an authentication error: {ex}",
         )
         exit(1)
+    except SmartDeviceException as ex:
+        _echo_error(
+            f"Unable to query {name} at once: {ex}",
+        )
+        if (
+            isinstance(ex, TimeoutException)
+            or ex.error_code == SmartErrorCode.SESSION_TIMEOUT_ERROR
+        ):
+            _echo_error(
+                "Timeout, try reducing the batch size via --batch-size option.",
+            )
+        exit(1)
     except Exception as ex:
-        click.echo(
-            click.style(f"Unable to query {name} at once: {ex}", bold=True, fg="red")
+        _echo_error(
+            f"Unexpected exception querying {name} at once: {ex}",
         )
         exit(1)
 
@@ -361,12 +382,8 @@ async def get_smart_fixture(device: TapoDevice, batch_size: int):
                 SmartRequest._create_request_dict(test_call.request)
             )
         except AuthenticationException as ex:
-            click.echo(
-                click.style(
-                    f"Unable to query the device due to an authentication error: {ex}",
-                    bold=True,
-                    fg="red",
-                )
+            _echo_error(
+                f"Unable to query the device due to an authentication error: {ex}",
             )
             exit(1)
         except Exception as ex: