From 824b6763e16b48d60da291323c736fa558883436 Mon Sep 17 00:00:00 2001
From: Ryan Summers <ryan.summers@vertigo-designs.com>
Date: Thu, 5 Aug 2021 16:21:55 +0200
Subject: [PATCH] Adding livestream HITL test

---
 hitl/run.sh       |   3 +
 hitl/streaming.py | 167 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 170 insertions(+)
 create mode 100644 hitl/streaming.py

diff --git a/hitl/run.sh b/hitl/run.sh
index 582c0f0..45acb03 100755
--- a/hitl/run.sh
+++ b/hitl/run.sh
@@ -33,3 +33,6 @@ python3 -m miniconf dt/sinara/dual-iir/04-91-62-d9-7e-5f afe/0='"G1"' iir_ch/0/0
 
 # Test the ADC/DACs connected via loopback.
 python3 hitl/loopback.py dt/sinara/dual-iir/04-91-62-d9-7e-5f
+
+# Test the livestream capabilities
+python3 hitl/streaming.py dt/sinara/dual-iir/04-91-62-d9-7e-5f
diff --git a/hitl/streaming.py b/hitl/streaming.py
new file mode 100644
index 0000000..1675409
--- /dev/null
+++ b/hitl/streaming.py
@@ -0,0 +1,167 @@
+#!/usr/bin/python3
+"""
+Author: Vertigo Designs, Ryan Summers
+
+Description: Implements HITL testing of Stabilizer data livestream capabilities.
+"""
+import asyncio
+import sys
+import time
+import argparse
+import socket
+import struct
+import logging
+
+from miniconf import Miniconf
+
+def _get_ip(broker):
+    """ Get the IP of the local device.
+
+    Args:
+        broker: The broker IP of the test. Used to select an interface to get the IP of.
+
+    Returns:
+        The IP as an array of integers.
+    """
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        sock.connect((broker, 1883))
+        address = sock.getsockname()[0]
+    finally:
+        sock.close()
+
+    return list(map(int, address.split('.')))
+
+
+def sequence_delta(previous_sequence, next_sequence):
+    """ Check the number of items between two sequence numbers. """
+    if previous_sequence is None:
+        return 0
+
+    delta = next_sequence - (previous_sequence + 1)
+    return delta & 0xFFFFFFFF
+
+
+class StabilizerStream:
+    """ Provides access to Stabilizer's livestreamed data. """
+
+    # The magic header half-word at the start of each packet.
+    MAGIC_HEADER = 0x057B
+
+    # The struct format of the header.
+    HEADER_FORMAT = '<HBBI'
+
+    # All supported formats by this reception script.
+    #
+    # The items in this dict are functions that will be provided the sample batch size and will
+    # return the struct deserialization code to unpack a single batch.
+    FORMAT = {
+        1: lambda batch_size: f'<{batch_size}H{batch_size}H{batch_size}H{batch_size}H'
+    }
+
+    def __init__(self, port):
+        """ Initialize the stream. """
+        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self.socket.bind(("", port))
+        self.socket.settimeout(0.5)
+
+
+    def clear(self, duration=5):
+        """ Clear the socket RX buffer by reading all available data.
+
+        Args:
+            duration: The maximum duration in seconds to read data for.
+        """
+        start = time.time()
+        while (time.time() - start) < duration:
+            try:
+                self.socket.recv(4096)
+            except socket.timeout:
+                return
+
+
+    def read_frame(self):
+        """ Read a single frame from the stream.
+
+        Returns:
+            Yields the (seqnum, data) of the batches available in the frame.
+        """
+        buf = self.socket.recv(4096)
+
+        # Attempt to parse a block from the buffer.
+        if len(buf) < struct.calcsize(self.HEADER_FORMAT):
+            return
+
+        # Parse out the packet header
+        magic, format_id, batch_size, sequence_number = struct.unpack_from(self.HEADER_FORMAT, buf)
+        buf = buf[struct.calcsize(self.HEADER_FORMAT):]
+
+        if magic != self.MAGIC_HEADER:
+            logging.warning('Encountered bad magic header: %s', hex(magic))
+            return
+
+        frame_format = self.FORMAT[format_id](batch_size)
+
+        batch_count = int(len(buf) / struct.calcsize(frame_format))
+
+        for offset in range(batch_count):
+            data = struct.unpack_from(frame_format, buf)
+            buf = buf[struct.calcsize(frame_format):]
+            yield (sequence_number + offset, data)
+
+
+def main():
+    """ Main program entry point. """
+    parser = argparse.ArgumentParser(description='Loopback tests for Stabilizer HITL testing',)
+    parser.add_argument('prefix', type=str,
+                        help='The MQTT topic prefix of the target')
+    parser.add_argument('--broker', '-b', default='mqtt', type=str,
+                        help='The MQTT broker address')
+    parser.add_argument('--port', '-p', default=2000, type=int,
+                        help='The UDP port to use for streaming')
+
+    args = parser.parse_args()
+
+    async def test():
+        """ The actual testing being completed. """
+        local_ip = _get_ip(args.broker)
+        interface = await Miniconf.create(args.prefix, args.broker)
+        stream = StabilizerStream(args.port)
+
+        # Configure the stream
+        print(f'Configuring stream to target {".".join(map(str, local_ip))}:{args.port}')
+        print('')
+        assert (await interface.command('stream_target',
+                                        {'ip': local_ip, 'port': args.port}))['code'] == 0
+
+        # Verify frame reception
+        print('Testing stream reception')
+        print('')
+        last_sequence = None
+        for _ in range(5000):
+            for (seqnum, _data) in stream.read_frame():
+                assert sequence_delta(last_sequence, seqnum) == 0
+                last_sequence = seqnum
+
+        # Disable the stream.
+        print('Closing stream')
+        print('')
+        assert (await interface.command('stream_target',
+                                        {'ip': [0, 0, 0, 0], 'port': 0}))['code'] == 0
+        stream.clear()
+
+        print('Verifying no further data is received')
+        try:
+            for _ in stream.read_frame():
+                raise Exception('Unexpected data encountered on stream')
+        except socket.timeout:
+            pass
+        print('PASS')
+
+
+    loop = asyncio.get_event_loop()
+    sys.exit(loop.run_until_complete(test()))
+
+
+if __name__ == '__main__':
+    main()