System Bridge Connector WebSocket Client Test

requirements.txt

systembridgeconnector>=5.0.0.dev2

systembridgeconnector_websocket_client_test.py

"""Test the WebSocket client for the System Bridge connector package."""

import asyncio
import logging
import sys
from dataclasses import asdict
from json import dumps
from typing import Final

from aiohttp import ClientSession
from systembridgeconnector._version import __version__
from systembridgeconnector.websocket_client import WebSocketClient
from systembridgemodels.modules import GetData, Module, ModulesData
from systembridgeshared.logger import setup_logger

API_HOST: Final[str] = "127.0.0.1"
API_PORT: Final[int] = 9170
API_TOKEN: Final[str] = sys.argv[1]

# Set up the logger
setup_logger("DEBUG", "websocket_client_test")
logger = logging.getLogger(__name__)


async def main() -> None:
    """Main function to test the WebSocket client."""

    # Log the version of the package
    logger.info("System Bridge connector: %s", __version__)

    # Create a new instance of the ClientSession
    client_session = ClientSession()
    logger.info("Created a new instance of the ClientSession.")

    # Create a new instance of the WebSocket client
    client = WebSocketClient(
        API_HOST,
        API_PORT,
        API_TOKEN,
        client_session,
    )
    logger.info("Created a new instance of the WebSocket client.")

    # Connect to the WebSocket server
    await client.connect()
    logger.info("Connected to the WebSocket server.")

    # Send a message to the WebSocket server
    module_data: ModulesData = await client.get_data(
        GetData(
            modules=[
                Module.BATTERY,
                Module.CPU,
                Module.DISKS,
                Module.DISPLAYS,
                Module.GPUS,
                Module.MEDIA,
                Module.MEMORY,
                Module.NETWORKS,
                Module.PROCESSES,
                Module.SENSORS,
                Module.SYSTEM,
            ]
        )
    )
    logger.info("Received module data: %s", module_data)

    # Close the WebSocket connection
    await client.close()
    logger.info("Closed the WebSocket connection.")

    # Close the ClientSession
    await client_session.close()
    logger.info("Closed the ClientSession.")

    # Save the output to a file
    with open(
        "output.json",
        "w",
        encoding="utf8",
    ) as file:
        file.write(dumps(asdict(module_data)))


if __name__ == "__main__":
    try:
        asyncio.new_event_loop().run_until_complete(main())
    except Exception:  # pylint: disable=broad-except
        logger.exception(
            "An error occurred while testing the WebSocket client.",
            exc_info=True,
        )
    except KeyboardInterrupt:
        pass