Merge pull request #108 from pyalarmdotcom/websocket-consolidated

Websocket consolidated

.devcontainer.json

@@ -24,7 +24,6 @@
                 "editor.formatOnPaste": false,
                 "editor.formatOnSave": true,
                 "editor.formatOnType": true,
-                // "mypy.enabled": true,
                 "python.analysis.autoSearchPaths": false,
                 "python.formatting.provider": "black",
                 "python.linting.enabled": true,

.vscode/settings.json

@@ -3,5 +3,6 @@
         "tests"
     ],
     "python.testing.unittestEnabled": false,
-    "python.testing.pytestEnabled": true
+    "python.testing.pytestEnabled": true,
+    "python.analysis.typeCheckingMode": "off"
 }

pyalarmdotcomajax/__init__.py

@@ -43,7 +43,7 @@
 # TODO: Use error handler and exception handlers in _async_get_system_devices on other request functions.
 # TODO: Fix get raw server response function.
 
-__version__ = "0.5.0-beta"
+__version__ = "0.5.0-beta.2"
 
 log = logging.getLogger(__name__)
 
@@ -164,6 +164,7 @@ def __init__(
         # CLI ATTRIBUTES
         #
         self.raw_catalog: dict = {}
+        self.raw_system: dict = {}
         self.raw_image_sensors: dict = {}
         self.raw_recent_images: dict = {}
 
@@ -349,7 +350,8 @@ async def async_update(self) -> None:  # noqa: C901
         #
 
         device_instances: dict[str, AllDevices_t] = {}
-        raw_devices: list[dict] = await self._async_get_system_devices(self._active_system_id)
+        raw_devices: list[dict] = await self._async_get_system(self._active_system_id)
+        raw_devices.extend(await self._async_get_system_devices(self._active_system_id))
 
         #
         # QUERY MULTI-DEVICE EXTENSIONS
@@ -377,7 +379,8 @@ async def async_update(self) -> None:  # noqa: C901
         #
         # BUILD PARTITIONS
         #
-        # Ensures that partition map is built before devices are built.
+
+        # Ensure that partition map is built before devices are built.
 
         for device in [
             device
@@ -633,15 +636,7 @@ def get_websocket_client(self) -> WebSocketClient:
 
         return WebSocketClient(self._websession, self._ajax_headers, self.devices)
 
-    #
-    #
-    #####################
-    # PRIVATE FUNCTIONS #
-    #####################
-    #
-    # Communicate directly with the ADC API
-
-    async def _is_logged_in(self) -> bool:
+    async def _async_keep_alive_login_check(self) -> bool:
         """Check if we are logged in."""
 
         async with self._websession.get(
@@ -652,6 +647,14 @@ async def _is_logged_in(self) -> bool:
 
         return bool(text_rsp == self.KEEP_ALIVE_CHECK_RESPONSE)
 
+    #
+    #
+    #####################
+    # PRIVATE FUNCTIONS #
+    #####################
+    #
+    # Communicate directly with the ADC API
+
     async def _async_get_active_system(self) -> str:
         """Get active system for user account."""
 
@@ -710,6 +713,8 @@ async def _async_has_image_sensors(self, system_id: str, retry_on_failure: bool
         Check is required because image sensors are not shown in the device catalog endpoint.
         """
 
+        # TODO: Needs changes to support multi-system environments
+
         try:
             log.info(f"Checking system {system_id} for image sensors.")
 
@@ -730,6 +735,32 @@ async def _async_has_image_sensors(self, system_id: str, retry_on_failure: bool
             log.error("Failed to get image sensors.")
             raise DataFetchFailed from err
 
+    async def _async_get_system(self, system_id: str, retry_on_failure: bool = True) -> list[dict]:
+        """Get all devices present in system."""
+
+        try:
+            log.info(f"Getting system data for {system_id}.")
+
+            async with self._websession.get(
+                url=AttributeRegistry.get_endpoints(DeviceType.SYSTEM)["primary"].format(c.URL_BASE, system_id),
+                headers=self._ajax_headers,
+                raise_for_status=True,
+            ) as resp:
+                json_rsp = await resp.json()
+
+                # Used by adc CLI.
+                self.raw_system = json_rsp
+
+                await self._async_handle_server_errors(json_rsp, "system", retry_on_failure)
+
+                return [json_rsp["data"]]
+
+        except (asyncio.TimeoutError, aiohttp.ClientError, aiohttp.ClientResponseError, KeyError) as err:
+            log.error("Failed to get system devices.")
+            raise DataFetchFailed from err
+        except TryAgain:
+            return await self._async_get_system(system_id=system_id, retry_on_failure=False)
+
     async def _async_get_system_devices(self, system_id: str, retry_on_failure: bool = True) -> list[dict]:
         """Get all devices present in system."""
 
@@ -946,7 +977,7 @@ async def _async_handle_server_errors(
                     )
                     raise DataFetchFailed(error_msg)
 
-                if not self._is_logged_in():
+                if not self._async_keep_alive_login_check():
                     log.debug(
                         "Error fetching data from Alarm.com. Got 403 status"
                         f" when requesting {request_name}. Trying to"

pyalarmdotcomajax/cli.py

@@ -9,6 +9,7 @@
 
 import argparse
 import asyncio
+import json
 import logging
 import platform
 import sys
@@ -261,11 +262,21 @@ async def cli() -> None:
             # Process MACHINE output
 
             device_type_output = {
-                slug_to_title(device_type.name): [
-                    device
-                    for device in alarm.raw_catalog.get("included", [])
-                    if device.get("type") == AttributeRegistry.get_relationship_id_from_devicetype(device_type)
-                ] or "\n(none found)\n"
+                slug_to_title(device_type.name): (
+                    json.dumps(filtered_list)
+                    if (
+                        filtered_list := [
+                            device
+                            for device in [
+                                *alarm.raw_catalog.get("included", []),
+                                alarm.raw_system.get("data"),
+                            ]
+                            if device.get("type")
+                            == AttributeRegistry.get_relationship_id_from_devicetype(device_type)
+                        ]
+                    )
+                    else "\n(none found)\n"
+                )
                 for device_type in DeviceType
                 if (
                     device_type in AttributeRegistry.supported_device_types
@@ -455,8 +466,10 @@ def _print_element_tearsheet(
         battery = "Critical"
     elif element.battery_low:
         battery = "Low"
-    else:
+    elif element.battery_critical is not None or element.battery_low is not None:
         battery = "Normal"
+    else:
+        battery = None
 
     # DESIRED STATE
     desired_str = (
@@ -468,30 +481,34 @@ def _print_element_tearsheet(
     # ATTRIBUTES
     output_str += "ATTRIBUTES: "
 
-    if isinstance(element.device_subtype, Sensor.Subtype):
-        output_str += f'[Type: {element.device_subtype.name.title().replace("_"," ")}] '
+    if isinstance(element.device_subtype, Sensor.Subtype) or element.state or battery or element.read_only:
+        if isinstance(element.device_subtype, Sensor.Subtype):
+            output_str += f'[TYPE: {element.device_subtype.name.title().replace("_"," ")}] '
 
-    if element.state:
-        output_str += f"[STATE: {element.state.name.title()}{desired_str}] "
+        if element.state:
+            output_str += f"[STATE: {element.state.name.title()}{desired_str}] "
 
-    output_str += f"[BATTERY: {battery}] "
+        if battery:
+            output_str += f"[BATTERY: {battery}] "
 
-    if element.read_only:
-        output_str += f"[READ ONLY: {element.read_only}] "
+        if element.read_only:
+            output_str += f"[READ ONLY: {element.read_only}] "
 
-    if isinstance(element, Light):
-        # Disabling. Boring stat.
-        # attribute_str += f"[REPORTS STATE: {element.supports_state_tracking}] "
+        if isinstance(element, Light):
+            # Disabling. Boring stat.
+            # attribute_str += f"[REPORTS STATE: {element.supports_state_tracking}] "
 
-        if element.brightness:
-            output_str += f"[BRIGHTNESS: {element.brightness}%] "
+            if element.brightness:
+                output_str += f"[BRIGHTNESS: {element.brightness}%] "
 
-    output_str += "\n"
+        # ENTITIES WITH "ATTRIBUTES" PROPERTY
+        if isinstance(element.attributes, BaseDevice.DeviceAttributes):
+            for name, value in asdict(element.attributes).items():
+                output_str += f"[{str(name).upper()}: {value}] "
+    else:
+        output_str += "(none)"
 
-    # ENTITIES WITH "ATTRIBUTES" PROPERTY
-    if isinstance(element.attributes, BaseDevice.DeviceAttributes):
-        for name, value in asdict(element.attributes).items():
-            output_str += f"[{str(name).upper()}: {value}] "
+        output_str += "\n"
 
     # SETTINGS / EXTENSIONS
 

pyalarmdotcomajax/websockets/client.py

@@ -52,12 +52,16 @@ class WebSocketClient:
     WEBSOCKET_ENDPOINT_TEMPLATE = "wss://webskt.alarm.com:8443/?auth={}"
     WEBSOCKET_TOKEN_REQUEST_TEMPLATE = "{}web/api/websockets/token"  # noqa: S105
 
-    def __init__(self, websession: ClientSession, ajax_headers: dict, device_registry: DeviceRegistry) -> None:
+    def __init__(
+        self,
+        websession: ClientSession,
+        ajax_headers: dict,
+        device_registry: DeviceRegistry,
+    ) -> None:
         """Initialize."""
         self._websession: ClientSession = websession
         self._ajax_headers: dict = ajax_headers
         self._device_registry: DeviceRegistry = device_registry
-
         self._ws_auth_token: str | None = None
 
     async def async_connect(self) -> None:
@@ -73,8 +77,7 @@ async def async_connect(self) -> None:
 
         # Connect to websocket endpoint
         async with self._websession.ws_connect(
-            self.WEBSOCKET_ENDPOINT_TEMPLATE.format(self._ws_auth_token),
-            headers=self._ajax_headers,
+            self.WEBSOCKET_ENDPOINT_TEMPLATE.format(self._ws_auth_token), headers=self._ajax_headers, timeout=30
         ) as websocket:
             async for msg in websocket:
                 # Message is JSON but encoded as text.

tests/conftest.py

@@ -85,6 +85,13 @@ def _load_mocks(repeat: bool = True) -> None:
             repeat=repeat,
         )
 
+        # response_mocker.get(
+        #     url=AttributeRegistry.get_endpoints(DeviceType.SYSTEM)["primary"].format(c.URL_BASE, ""),
+        #     status=200,
+        #     body=get_http_body_json("systems_ok"),
+        #     repeat=repeat,
+        # )
+
         response_mocker.get(
             url=AttributeRegistry.get_endpoints(DeviceType.IMAGE_SENSOR)["primary"].format(c.URL_BASE, ""),
             status=200,