python-kasa · sdb9696 · Feb 20, 2024 · Feb 19, 2024 · Feb 19, 2024 · Feb 19, 2024
diff --git a/kasa/cli.py b/kasa/cli.py
@@ -90,19 +90,43 @@ def wrapper(message=None, *args, **kwargs):
 pass_dev = click.make_pass_decorator(Device)
 
 
-class ExceptionHandlerGroup(click.Group):
-    """Group to capture all exceptions and print them nicely.
+def CatchAllExceptions(cls):
+    """Capture all exceptions and prints them nicely.
 
-    Idea from https://stackoverflow.com/a/44347763
+    Idea from https://stackoverflow.com/a/44347763 and
+    https://stackoverflow.com/questions/52213375
     """
 
-    def __call__(self, *args, **kwargs):
-        """Run the coroutine in the event loop and print any exceptions."""
-        try:
-            asyncio.get_event_loop().run_until_complete(self.main(*args, **kwargs))
-        except Exception as ex:
-            echo(f"Got error: {ex!r}")
+    def _handle_exception(debug, exc):
+        if isinstance(exc, click.ClickException):
+            raise
+        echo(f"Raised error: {exc}")
+        if debug:
             raise
+        echo("Run with --debug enabled to see stacktrace")
+        sys.exit(1)
+
+    class _CommandCls(cls):
+        _debug = False
+
+        async def make_context(self, info_name, args, parent=None, **extra):
+            self._debug = any(
+                [arg for arg in args if arg in ["--debug", "-d", "--verbose", "-v"]]
+            )
+            try:
+                return await super().make_context(
+                    info_name, args, parent=parent, **extra
+                )
+            except Exception as exc:
+                _handle_exception(self._debug, exc)
+
+        async def invoke(self, ctx):
+            try:
+                return await super().invoke(ctx)
+            except Exception as exc:
+                _handle_exception(self._debug, exc)
+
+    return _CommandCls
 
 
 def json_formatter_cb(result, **kwargs):
@@ -129,7 +153,7 @@ def _device_to_serializable(val: Device):
 
 @click.group(
     invoke_without_command=True,
-    cls=ExceptionHandlerGroup,
+    cls=CatchAllExceptions(click.Group),
     result_callback=json_formatter_cb,
 )
 @click.option(

diff --git a/kasa/tests/test_cli.py b/kasa/tests/test_cli.py
@@ -504,6 +504,7 @@ async def test_host_unsupported(unsupported_device_info):
             "foo",
             "--password",
             "bar",
+            "--debug",
         ],
     )
 
@@ -563,6 +564,7 @@ async def test_host_auth_failed(discovery_mock, mocker):
             "foo",
             "--password",
             "bar",
+            "--debug",
         ],
     )
 
@@ -610,3 +612,49 @@ async def test_shell(dev: Device, mocker):
     res = await runner.invoke(cli, ["shell"], obj=dev)
     assert res.exit_code == 0
     embed.assert_called()
+
+
+async def test_errors(mocker):
+    runner = CliRunner()
+    err = SmartDeviceException("Foobar")
+
+    # Test masking
+    mocker.patch("kasa.Discover.discover", side_effect=err)
+    res = await runner.invoke(
+        cli,
+        ["--username", "foo", "--password", "bar"],
+    )
+    assert res.exit_code == 1
+    assert "Raised error: Foobar" in res.output
+    assert "Run with --debug enabled to see stacktrace" in res.output
+    assert isinstance(res.exception, SystemExit)
+
+    # Test --debug
+    res = await runner.invoke(
+        cli,
+        ["--debug"],
+    )
+    assert res.exit_code == 1
+    assert "Raised error: Foobar" in res.output
+    assert res.exception == err
+
+    # Test no device passed to subcommand
+    mocker.patch("kasa.Discover.discover", return_value={})
+    res = await runner.invoke(
+        cli,
+        ["sysinfo"],
+    )
+    assert res.exit_code == 1
+    assert (
+        "Raised error: Managed to invoke callback without a context object of type 'Device' existing."
+        in res.output
+    )
+    assert isinstance(res.exception, SystemExit)
+
+    # Test click error
+    res = await runner.invoke(
+        cli,
+        ["--foobar"],
+    )
+    assert res.exit_code == 2
+    assert "Raised error:" not in res.output