mosquitto/test/mosquitto_broker.py

"""
"""

from subprocess import Popen
from typing import Optional, Union

import logging
import os

import mosq_test

from broker_config import BrokerConfig

class MosquittoBroker:
    def __init__(
        self,
        port: Optional[int] = None,
        config: Union[None, BrokerConfig, str, list, dict] = None,
        config_file_name=None,
        add_port_to_config=True,
        remove_files=True,
        env=None,
        termination_timeout = 10,
        expect_fail=False,
        check_port=True,
    ):
        assert port is not None or (isinstance(config, BrokerConfig))
        if isinstance(config, BrokerConfig):
            self._port = config.listeners[0].port if len(config.listeners) > 0 else port
        else:
            self._port = port

        self._process :Optional[Popen]= None
        self.env = env
        self.check_port = check_port
        if config_file_name is None or config_file_name == "":
            config_file_name = f"{str(self._port)}.conf"
        self._config_file_name = config_file_name
        self._config_files = set()
        self._expect_fail = expect_fail
        self._extra_files = set()
        self._remove_files = remove_files
        self._termination_timeout = termination_timeout
        if config:
            if (
                not isinstance(config, BrokerConfig)
                and add_port_to_config
                and self._port
            ):
                self.add_config(f"listener {self._port}\n")
            self.add_config(config)

    def __str__(self):
        return f"{self.__class__.__name__}:{self._port}"

    @property
    def port(self):
        return self._port

    @property
    def process(self):
        if self._process is None:
            raise RuntimeError("process not started yet")
        return self._process

    def set_remove_files(self, value: bool):
        self._remove_files = value

    def add_config(self, config, file_name=None):
        if isinstance(config, dict):
            config = "\n".join(f"{k} {v}" for k, v in config.items())
        if isinstance(config, list):
            config = "\n".join(f"{k} {v}" for k, v in config)
        if isinstance(config, BrokerConfig):
            config = str(config)
        if file_name is None or file_name == "":
            file_name = self._config_file_name
        file_mode = "w"
        if file_name in self._config_files:
            file_mode = "a"
        else:
            self._config_files.add(file_name)
        logging.debug(f"{self} Write to config file {file_name} with mode {file_mode}")
        with open(file_name, file_mode) as config_file:
            config_file.writelines(config)

    def add_extra_file(self, file_name):
        self._extra_files.add(file_name)

    def start(self, expect_fail=False, start_timeout=0.1):
        logging.info(f"Starting {self}")
        self._process = mosq_test.start_broker(
            use_conf=len(self._config_files) > 0,
            filename=self._config_file_name,
            port=self._port,
            env=self.env,
            expect_fail=expect_fail,
            check_port=self.check_port,
            timeout=start_timeout
        )
        logging.info(f"{self} "+ "Started" if self.is_running() else f"Terminated rc={self._process.returncode}")

    def is_running(self):
        return self._process and self._process.poll() is None

    def check_log(self, matcher):
        if not matcher(self.get_log()):
            raise AssertionError(f"{matcher.last_message}")
        else:
            logging.info(f"{self} Check log {matcher}")

    def get_log(self):
        return mosq_test.broker_log(self._process)

    def terminate(self):
        if self._process:
            mosq_test.terminate_broker(self._process)

    def reload(self):
        if self._process:
            mosq_test.reload_broker(self._process)

    def __enter__(self):
        if self._process == None:
            self.start()
        return self

    def __exit__(self, ex_type, value, tb):
        self.stop(ex_type)

    def stop(self, ex_type=None):
        if self._process:
            timed_out, _ = mosq_test.terminate_broker(self._process)
            logging.info(f"Stopping {self}")
            logging.info(f"Stopped {self}")
            if ex_type is not None or timed_out or (self._expect_fail == False and self._process.returncode != 0):
                print(f"\n{self} log:")
                print(self.get_log())
            else:
                logging.debug(f"\n{self} log:")
                logging.debug(self.get_log())
            if timed_out:
                raise RuntimeError(f"{self} timed out when shutting down")
            if self._expect_fail == False and self._process.returncode != 0:
                raise RuntimeError(f"{self} exited with {self._process.returncode}")

        if self._remove_files:
            for file_name in self._config_files:
                os.remove(file_name)
            for file_name in self._extra_files:
                try:
                    os.remove(file_name)
                except FileNotFoundError:
                    pass