DynaLab/test/minimal_connector.py

# minimal_connector.py
from __future__ import annotations

import asyncio
from random import random
from time import monotonic
import uuid
from dynalab_protocol.models import (
    ConnectorHello,
    DynaLabHello,
    HandshakeAccepted,
    HandshakeRejected,
    Heartbeat,
    SignalDescriptor,
    ValueDescriptor,
    machine_timestamp_ms,
)
from dynalab_protocol.wire import read_message, write_message


HOST = "127.0.0.1"
PORT = 8765
RETRY_DELAY_SECONDS = 2


def build_signals() -> list[SignalDescriptor]:
    """Describe the signals this example connector exposes to DynaLab."""
    return [
        SignalDescriptor(
            id=uuid.uuid4(), name="lambda", type="number", min_value=0, max_value=5
        ),
        SignalDescriptor(
            id=uuid.uuid4(),
            name="torque",
            type="number",
            min_value=0,
            max_value=10,
            unit="Nm",
        ),
        SignalDescriptor(
            id=uuid.uuid4(),
            name="RPM",
            type="number",
            min_value=0,
            max_value=6000,
            unit="RPM",
        ),
        SignalDescriptor(
            id=uuid.uuid4(),
            name="temperature",
            type="number",
            min_value=0,
            max_value=120,
            unit="°C",
        ),
        SignalDescriptor(
            id=uuid.uuid4(),
            name="voltage",
            type="number",
            min_value=0,
            max_value=20,
            unit="V",
        ),
        SignalDescriptor(
            id=uuid.uuid4(),
            name="current",
            type="number",
            min_value=0,
            max_value=50,
            unit="A",
        ),
    ]


async def _value_sender(
    writer: asyncio.StreamWriter, connector_hello: ConnectorHello
) -> None:
    last_value_send = monotonic()
    while True:
        if monotonic() > last_value_send + 0.1:
            multipliers: list[int] = [5, 10, 6000, 120, 20, 50]
            for i in range(6):
                random_val = random() * multipliers[i]
                value = ValueDescriptor(
                    signal_id=connector_hello.signals[i].id, value=random_val
                )
                await write_message(writer, value)
            last_value_send = monotonic()

        else:
            await asyncio.sleep(0.05)


async def run_connector_session() -> None:
    """Open one connection to DynaLab and run the connector protocol."""
    reader, writer = await asyncio.open_connection(HOST, PORT)

    dynalab_hello: DynaLabHello
    handshake: HandshakeAccepted

    peer = writer.get_extra_info("peername")
    print(f"Connected to DynaLab at {peer}")

    try:
        # DynaLab starts the protocol by introducing itself to the connector.
        message = await asyncio.wait_for(read_message(reader), timeout=5.0)

        if not isinstance(message, DynaLabHello):
            raise ValueError(f"Expected dynalab_hello, received {message.type}")

        dynalab_hello = message
        print(f"Received DynaLab Hello\n {dynalab_hello}")

        # Reply with the connector identity and signal list.
        connector_hello = ConnectorHello(
            connector_uuid=uuid.uuid4(),
            connector_name="Minimal Connector",
            connector_version="a0.0.1",
            signals=build_signals(),
        )
        await write_message(writer, connector_hello)

        # DynaLab accepts or rejects the connector after validating the hello.
        message = await asyncio.wait_for(read_message(reader), timeout=5.0)

        if isinstance(message, HandshakeAccepted):
            handshake = message
            print(f"Received connection handshake: {handshake}")
        elif isinstance(message, HandshakeRejected):
            print(f"Connector handshake rejected: {message.reason}")
            return

        value_task = asyncio.create_task(_value_sender(writer, connector_hello))
        # Keep the connection alive by returning heartbeat messages.
        while True:
            message = await read_message(reader)
            print(f"Received {message}")

            if isinstance(message, Heartbeat):
                if message.sequence < 10:
                    message.return_timestamp = machine_timestamp_ms()
                    await write_message(writer, message)

    except asyncio.CancelledError:
        raise

    finally:
        value_task.cancel()
        print("Disconnecting")
        writer.close()
        await writer.wait_closed()


async def main() -> None:
    while True:
        try:
            await run_connector_session()
            return
        except ConnectionError as error:
            print(
                f"Connection error: {error}. "
                f"Retrying in {RETRY_DELAY_SECONDS} seconds..."
            )
            await asyncio.sleep(RETRY_DELAY_SECONDS)


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nConnector stopped")