sf_daq_broker/tests/test_broker_manager.py

import json
import unittest

import os
from time import time

from sf_databuffer_writer import config
from sf_databuffer_writer.broker_manager import BrokerManager
from sf_databuffer_writer.utils import verify_channels


class MockRequestSender(object):
    def __init__(self):
        self.write_request = None
        self.sendto_epics_writer = None

    def send(self, write_request, sendto_epics_writer):
        self.write_request = write_request
        self.sendto_epics_writer = sendto_epics_writer


class TestBrokerManager(unittest.TestCase):
    TEST_AUDIT_FILE = "ignore_output.txt"

    def tearDown(self):
        try:
            os.remove(TestBrokerManager.TEST_AUDIT_FILE)
        except:
            pass

    def test_write_request(self):

        request_sender = MockRequestSender()
        channels = ["test_1", "test_2"]
        start_pulse_id = 100
        stop_pulse_id = 120
        parameters = {"general/created": "test",
                      "general/user": "tester",
                      "general/process": "test_process",
                      "general/instrument": "mac",
                      "output_file": "test.h5"}

        manager = BrokerManager(request_sender, channels, TestBrokerManager.TEST_AUDIT_FILE)

        manager.set_parameters(parameters)
        manager.start_writer(start_pulse_id)
        manager.stop_writer(stop_pulse_id)

        data_api_request = json.loads(request_sender.write_request["data_api_request"])
        request_parameters = json.loads(request_sender.write_request["parameters"])

        self.assertEqual(data_api_request["range"]["startPulseId"], start_pulse_id)
        self.assertEqual(data_api_request["range"]["endPulseId"], stop_pulse_id)
        self.assertListEqual(data_api_request["channels"], [{'name': ch, 'backend': config.DATA_BACKEND}
                                                            for ch in channels])

        self.assertDictEqual(request_parameters, parameters)

        self.assertListEqual(data_api_request["eventFields"], ["channel", "pulseId", "value", "shape", "globalDate"])
        self.assertListEqual(data_api_request["configFields"], ["type", "shape"])

        start_pulse_id = 1000
        stop_pulse_id = 1100
        parameters = {"general/created": "test2",
                      "general/user": "tester2",
                      "general/process": "test_process2",
                      "general/instrument": "mac2",
                      "output_file": "test2.h5"}

        manager.set_parameters(parameters)
        manager.start_writer(start_pulse_id)
        manager.stop_writer(stop_pulse_id)

        data_api_request = json.loads(request_sender.write_request["data_api_request"])
        request_parameters = json.loads(request_sender.write_request["parameters"])
        request_timestamp = request_sender.write_request["timestamp"]

        self.assertEqual(data_api_request["range"]["startPulseId"], start_pulse_id)
        self.assertEqual(data_api_request["range"]["endPulseId"], stop_pulse_id)

        self.assertDictEqual(request_parameters, parameters)

        self.assertTrue(time() - request_timestamp < 1)
        self.assertTrue(request_sender.sendto_epics_writer)

    def test_audit_file(self):
        request_sender = MockRequestSender()
        config.SEPARATE_CAMERA_CHANNELS = False

        channels = ["test_1", "test_2"]
        start_pulse_id = 100
        stop_pulse_id = 120
        parameters = {"general/created": "test",
                      "general/user": "tester",
                      "general/process": "test_process",
                      "general/instrument": "mac",
                      "output_file": "test.h5"}

        manager = BrokerManager(request_sender, channels, TestBrokerManager.TEST_AUDIT_FILE)

        manager.set_parameters(parameters)
        manager.start_writer(start_pulse_id)
        manager.stop_writer(stop_pulse_id)

        data_api_request = json.loads(request_sender.write_request["data_api_request"])
        request_parameters = json.loads(request_sender.write_request["parameters"])

        with open(TestBrokerManager.TEST_AUDIT_FILE) as input_file:
            lines = input_file.readlines()

        audit_log = json.loads(lines[0][18:])
        audit_data_api_request = json.loads(audit_log["data_api_request"])
        audit_request_parameters = json.loads(audit_log["parameters"])
        audit_timestamp = audit_log["timestamp"]

        self.assertDictEqual(data_api_request, audit_data_api_request)
        self.assertDictEqual(request_parameters, audit_request_parameters)

        self.assertTrue(time()-audit_timestamp < 1)

    def test_manager_status(self):
        request_sender = MockRequestSender()
        channels = []

        manager = BrokerManager(request_sender, channels, TestBrokerManager.TEST_AUDIT_FILE)

        self.assertEqual(manager.get_status(), "stopped")

        parameters = {"general/created": "test",
                      "general/user": "tester",
                      "general/process": "test_process",
                      "general/instrument": "mac",
                      "output_file": "test.h5"}

        manager.set_parameters(parameters)

        self.assertEqual(manager.get_status(), "configured")

        manager.start_writer(100)

        self.assertEqual(manager.get_status(), "receiving")

        manager.stop_writer(120)

        self.assertEqual(manager.get_status(), "stopped")

    def test_verify_channels(self):
        config.BROKER_CHANNELS_LIMIT = 4
        config.BROKER_CHANNELS_LIMIT_PICTURE = 1

        channels = ["test1", "test2", "test3", "test4"]
        verify_channels(channels)

        channels = ["test1", "test2", "test3", "test4:FPICTURE", ""]
        verify_channels(channels)

        channels = ["test1", "test2", "test3", "test4:FPICTURE", "", ""]
        verify_channels(channels)

        channels = ["test1", "test2", "test3", "test4", "test5"]
        with self.assertRaisesRegex(ValueError, "Too many bsread channels"):
            verify_channels(channels)

        channels = ["test1", "test2", "test3", "test4", "test5:FPICTURE"]
        with self.assertRaisesRegex(ValueError, "Too many bsread channels"):
            verify_channels(channels)

        channels = ["test1", "test2", "test3:FPICTURE", "test4:FPICTURE"]
        with self.assertRaisesRegex(ValueError, "Too many picture channels"):
            verify_channels(channels)