Allow for configurable larger websocket responses

homeassistant_cli/cli.py

@@ -24,6 +24,18 @@
 )
 
 
+def validate_file_size(ctx, param, value):
+    if value is not None:
+        if str(value).endswith('KB'):
+            size = int(value[:-2]) * 1024
+        elif str(value).endswith('MB'):
+            size = int(value[:-2]) * 1024 * 1024
+        else:
+            # bytes
+            size = int(value)
+        return size
+
+
 def run() -> None:
     """Run entry point.
 
@@ -136,7 +148,7 @@ def _default_token() -> Optional[str]:
 @click.option(
     '--timeout',
     help='Timeout for network operations.',
-    default=const.DEFAULT_TIMEOUT,
+    default=str(const.DEFAULT_TIMEOUT),
     show_default=True,
 )
 @click.option(
@@ -205,6 +217,13 @@ def _default_token() -> Optional[str]:
     default=None,
     help='Sort table by the jsonpath expression. Example: last_changed',
 )
+@click.option(
+    '--max-message-size',
+    callback=validate_file_size,
+    default=const.WS_MAX_MESSAGE_SIZE,
+    help='Max size of websocket payload. Default: 4MB',
+    envvar='HASS_WS_MAX_MESSAGE_SIZE',
+)
 @pass_context
 def cli(
     ctx: Configuration,
@@ -222,6 +241,7 @@ def cli(
     no_headers: bool,
     table_format: str,
     sort_by: Optional[str],
+    max_message_size: str,
 ) -> None:
     """Command line interface for Home Assistant."""
     ctx.verbose = verbose
@@ -238,6 +258,7 @@ def cli(
     ctx.no_headers = no_headers
     ctx.table_format = table_format
     ctx.sort_by = sort_by  # type: ignore
+    ctx.max_message_size = max_message_size
 
     _LOGGER.debug("Using settings: %s", ctx)
 

homeassistant_cli/config.py

@@ -84,9 +84,7 @@ def resolve_server(ctx: Any) -> str:  # noqa: F821
         ctx.resolved_server = None
 
     if not ctx.resolved_server:
-
         if ctx.server == "auto":
-
             if "HASSIO_TOKEN" in os.environ and "HASS_TOKEN" not in os.environ:
                 ctx.resolved_server = const.DEFAULT_SERVER_MDNS
             else:
@@ -126,6 +124,7 @@ def __init__(self) -> None:
         self.no_headers = False
         self.table_format = 'plain'
         self.sort_by = None
+        self.max_message_size = const.WS_MAX_MESSAGE_SIZE  # type: str
 
     def echo(self, msg: str, *args: Optional[Any]) -> None:
         """Put content message to stdout."""
@@ -151,6 +150,7 @@ def __repr__(self) -> str:
             "access-token": 'yes' if self.token is not None else 'no',
             "api-password": 'yes' if self.password is not None else 'no',
             "insecure": self.insecure,
+            "max_message_size": self.max_message_size,
             "output": self.output,
             "verbose": self.verbose,
         }

homeassistant_cli/const.py

@@ -20,3 +20,5 @@
     ('CHANGED', 'last_changed'),
 ]
 COLUMNS_SERVICES = [('DOMAIN', 'domain'), ("SERVICE", "domain.services[*]")]
+
+WS_MAX_MESSAGE_SIZE = '4194304'

homeassistant_cli/remote.py

@@ -105,9 +105,9 @@ def wsapi(
     async def fetcher() -> Optional[Dict]:
         async with aiohttp.ClientSession() as session:
             async with session.ws_connect(
-                resolve_server(ctx) + "/api/websocket"
+                resolve_server(ctx) + "/api/websocket",
+                max_msg_size=ctx.max_message_size,
             ) as wsconn:
-
                 await wsconn.send_str(
                     json.dumps({'type': 'auth', 'access_token': ctx.token})
                 )

tests/test_area.py

@@ -13,7 +13,6 @@ def test_area_list(default_areas) -> None:
     with mock.patch(
         'homeassistant_cli.remote.get_areas', return_value=default_areas
     ):
-
         runner = CliRunner()
         result = runner.invoke(
             cli.cli, ["--output=json", "area", "list"], catch_exceptions=False
@@ -29,7 +28,6 @@ def test_area_list_filter(default_areas) -> None:
     with mock.patch(
         'homeassistant_cli.remote.get_areas', return_value=default_areas
     ):
-
         runner = CliRunner()
         result = runner.invoke(
             cli.cli,

tests/test_device.py

@@ -16,7 +16,6 @@ def test_device_list(default_devices, default_areas) -> None:
         with mock.patch(
             'homeassistant_cli.remote.get_areas', return_value=default_areas
         ):
-
             runner = CliRunner()
             result = runner.invoke(
                 cli.cli,
@@ -37,7 +36,6 @@ def test_device_list_filter(default_devices, default_areas) -> None:
         with mock.patch(
             'homeassistant_cli.remote.get_areas', return_value=default_areas
         ):
-
             runner = CliRunner()
             result = runner.invoke(
                 cli.cli,
@@ -64,7 +62,6 @@ def test_device_assign(default_areas, default_devices) -> None:
                 'homeassistant_cli.remote.assign_area',
                 return_value={'success': True},
             ):
-
                 runner = CliRunner()
                 result = runner.invoke(
                     cli.cli,

tests/test_raw.py

@@ -82,7 +82,6 @@ def test_raw_ws() -> None:
     with mocker.patch(
         'homeassistant_cli.remote.wsapi', return_value={"result": "worked"}
     ) as mockmethod:
-
         runner = CliRunner()
         result = runner.invoke(
             cli.cli,
@@ -103,7 +102,6 @@ def test_raw_ws_data() -> None:
     with mocker.patch(
         'homeassistant_cli.remote.wsapi', return_value={"result": "worked"}
     ) as mockmethod:
-
         runner = CliRunner()
         result = runner.invoke(
             cli.cli,

tests/test_service.py

@@ -74,7 +74,6 @@ def test_service_completion(default_services) -> None:
 def test_service_call(default_services) -> None:
     """Test basic call of a service."""
     with requests_mock.Mocker() as mock:
-
         post = mock.post(
             "http://localhost:8123/api/services/homeassistant/restart",
             json={"result": "bogus"},