From f96b6682429b035a2a46242b07d715fecea8d639 Mon Sep 17 00:00:00 2001 From: Ilia Sotnikov Date: Tue, 10 Sep 2024 01:20:06 +0300 Subject: [PATCH] feat: Added support for dry-run mode and overrides for debug/oneshot from command line * Dry-run mode, when enabled with `-a` command line switch, will skip all MQTT calls, including connection and disconnection, as well as publishing messages. This is useful for testing the application without actually sending any data over MQTT * One-shot and debug modes could now be enabled from command line with `-o` and `-d` switches respectively --- src/energomera_hass_mqtt/hass_mqtt.py | 4 ++- src/energomera_hass_mqtt/main.py | 32 ++++++++++++++++++++-- src/energomera_hass_mqtt/mqtt_client.py | 36 ++++++++++++++++++++++++- tests/conftest.py | 17 ++++++++++++ tests/test_energomera.py | 36 ++++++++++++++++++++++++- 5 files changed, 120 insertions(+), 5 deletions(-) diff --git a/src/energomera_hass_mqtt/hass_mqtt.py b/src/energomera_hass_mqtt/hass_mqtt.py index 9ca6280..a5771fc 100644 --- a/src/energomera_hass_mqtt/hass_mqtt.py +++ b/src/energomera_hass_mqtt/hass_mqtt.py @@ -74,9 +74,10 @@ def calculate_bcc(bytes_data: bytes) -> bytes: return bcc.to_bytes(length=1, byteorder='big') def __init__( - self, config: EnergomeraConfig, + self, config: EnergomeraConfig, dry_run: bool = False ): self._config = config + self._dry_run = dry_run # Override the method in the `iec62056_21` library with the specific # implementation # pylint: disable=protected-access @@ -100,6 +101,7 @@ def __init__( mqtt_tls_context = ssl.SSLContext() self._mqtt_client = MqttClient( + dry_run=self._dry_run, hostname=getenv('MQTT_HOST', config.of.mqtt.host), port=getenv('MQTT_PORT', config.of.mqtt.port), username=getenv('MQTT_USER', config.of.mqtt.user), diff --git a/src/energomera_hass_mqtt/main.py b/src/energomera_hass_mqtt/main.py index b941c20..64e3690 100644 --- a/src/energomera_hass_mqtt/main.py +++ b/src/energomera_hass_mqtt/main.py @@ -49,6 +49,24 @@ def process_cmdline() -> Tuple[str, Namespace]: default=DEFAULT_CONFIG_FILE, help="Path to configuration file (default: '%(default)s')" ) + parser.add_argument( + '-a', '--dry-run', + action='store_true', + default=False, + help="Dry run, do not actually send any data" + ) + parser.add_argument( + '-d', '--debug', + action='store_true', + default=False, + help="Enable debug logging" + ) + parser.add_argument( + '-o', '--one-shot', + action='store_true', + default=False, + help="Run only once, then exit" + ) return (parser.prog, parser.parse_args()) @@ -70,13 +88,23 @@ async def async_main(mqtt_port: Optional[int] = None) -> None: # Override port of MQTT broker (if provided) if mqtt_port: config.of.mqtt.port = mqtt_port - logging.basicConfig(level=config.logging_level) + + # Support overriding certain configuration options from command line + if args.one_shot: + config.of.general.oneshot = True + + if args.debug: + config.of.general.logging_level = 'debug' # Print configuration details print(f'Starting {prog}, configuration:') print(config) - client = EnergomeraHassMqtt(config) + if args.dry_run: + print('Dry run mode enabled, no data will be sent over MQTT') + + logging.basicConfig(level=config.logging_level) + client = EnergomeraHassMqtt(config=config, dry_run=args.dry_run) while True: config.interpolate() try: diff --git a/src/energomera_hass_mqtt/mqtt_client.py b/src/energomera_hass_mqtt/mqtt_client.py index 84c6d99..5b04fcc 100644 --- a/src/energomera_hass_mqtt/mqtt_client.py +++ b/src/energomera_hass_mqtt/mqtt_client.py @@ -60,11 +60,15 @@ class MqttClient(aiomqtt.Client): last will thru its constructor, while the `EnergomeraHassMqtt` consuming it only has required property for that around actual publish calls. + :param dry_run: Whether to skip actual connection and publish calls :param args: Pass-through positional arguments for parent class :param kwargs: Pass-through keyword arguments for parent class """ - def __init__(self, *args: Any, **kwargs: Any) -> None: + def __init__( + self, *args: Any, dry_run: bool = False, **kwargs: Any + ) -> None: self._will_set = 'will' in kwargs + self._dry_run = dry_run super().__init__(logger=_LOGGER, *args, **kwargs) # Skip locking in `__aenter__` and `__aexit__` - those aren't used with # context manager, rather as regular methods, especially the former @@ -81,6 +85,9 @@ async def connect(self) -> None: a process loop with no risk of hitting non-reentrant error from base class. """ + if self._dry_run: + _LOGGER.debug('Dry run, skipping connection') + return # Using combination of `self._connected` and `self._disconnected` (both # inherited from `asyncio_mqtt.client`) to detect of MQTT client needs # a reconnection upon a network error isn't reliable - the former isn't @@ -104,9 +111,27 @@ async def disconnect(self) -> None: Disconnects from MQTT broker using `__aexit__` method of the base class as per migration guide above. """ + if self._dry_run: + _LOGGER.debug('Dry run, skipping disconnection') + return + # pylint:disable=unnecessary-dunder-call await self.__aexit__(None, None, None) + # pylint:disable=arguments-differ + async def publish(self, *args: Any, **kwargs: Any) -> None: + """ + Publishes a message to the MQTT broker. + + :param args: Pass-through positional arguments for parent method + :param kwargs: Pass-through keyword arguments for parent method + """ + if self._dry_run: + _LOGGER.debug('Dry run, skipping publish') + return + + await super().publish(*args, **kwargs) + def will_set(self, *args: Any, **kwargs: Any) -> None: """ Allows setting last will to the underlying MQTT client. @@ -117,8 +142,13 @@ def will_set(self, *args: Any, **kwargs: Any) -> None: :param args: Pass-through positional arguments for parent method :param kwargs: Pass-through keyword arguments for parent method """ + if self._dry_run: + _LOGGER.debug('Dry run, skipping to set the last will') + return + if self._will_set: return + self._client.will_set(*args, **kwargs) self._will_set = True @@ -126,5 +156,9 @@ def will_clear(self) -> None: """ Clears the last will might have been set previously. """ + if self._dry_run: + _LOGGER.debug('Dry run, skipping to clear the last will') + return + self._client.will_clear() self._will_set = False diff --git a/tests/conftest.py b/tests/conftest.py index 0339af5..cded375 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -28,6 +28,7 @@ import sys from functools import reduce from unittest.mock import patch, call, DEFAULT, AsyncMock, Mock +import aiomqtt from callee import Regex as CallRegexMatcher import pytest from pytest import FixtureRequest @@ -1104,3 +1105,19 @@ def mock_mqtt() -> Iterator[MockMqttT]: publish=DEFAULT, connect=DEFAULT, will_set=DEFAULT, new_callable=AsyncMock) as mocks: yield mocks + + +@pytest.fixture +def mock_mqtt_underlying() -> Iterator[MockMqttT]: + ''' + Provides necessary mocks to underlying MQTT client, to be used as context + manager. + ''' + with patch.multiple( + aiomqtt.Client, + publish=DEFAULT, + __aenter__=DEFAULT, + __aexit__=DEFAULT, + new_callable=AsyncMock + ) as mocks: + yield mocks diff --git a/tests/test_energomera.py b/tests/test_energomera.py index c2cdd83..e1537d8 100644 --- a/tests/test_energomera.py +++ b/tests/test_energomera.py @@ -21,7 +21,8 @@ ''' Tests for 'energomera_hass_mqtt' package ''' -from unittest.mock import call +from unittest.mock import call, patch +import logging import pytest from conftest import ( MQTT_PUBLISH_CALLS_COMPLETE, SERIAL_EXCHANGE_COMPLETE, MockMqttT, @@ -42,6 +43,39 @@ def test_normal_run(mock_serial: MockSerialT, mock_mqtt: MockMqttT) -> None: ) +@pytest.mark.usefixtures('mock_config') +def test_dry_run( + mock_serial: MockSerialT, mock_mqtt_underlying: MockMqttT, + monkeypatch: pytest.MonkeyPatch +) -> None: + ''' + Tests for enabling dry-run mode from command line. + ''' + monkeypatch.setattr('sys.argv', ['energomera_hass_mqtt', '-a']) + main() + # Ensure that no MQTT calls are made are sent + mock_mqtt_underlying['publish'].assert_not_called() + mock_mqtt_underlying['__aenter__'].assert_not_called() + mock_mqtt_underlying['__aexit__'].assert_not_called() + # While serial exchanges are still made + mock_serial['_send'].assert_has_calls( + [call(x['receive_bytes']) for x in SERIAL_EXCHANGE_COMPLETE] + ) + + +@pytest.mark.usefixtures('mock_config', 'mock_mqtt', 'mock_serial') +def test_debug_run( + monkeypatch: pytest.MonkeyPatch +) -> None: + ''' + Tests for enabling debug logging from command line. + ''' + monkeypatch.setattr('sys.argv', ['energomera_hass_mqtt', '-d']) + with patch('logging.basicConfig') as mock: + main() + mock.assert_called_with(level=logging.DEBUG) + + @pytest.mark.usefixtures('mock_serial', 'mock_config', 'mock_mqtt') @pytest.mark.serial_simulate_timeout(True) def test_timeout() -> None: