measure_ctrl/tests/test_mqtt_service.py

import pytest
from unittest.mock import MagicMock, patch, call
import json
from src.services.mqtt_service import MQTTService, InsertedCell
from src.models.device import DeviceStatus

@pytest.fixture
def mock_mqtt_client():
    with patch('paho.mqtt.client.Client') as mock:
        client_instance = MagicMock()
        mock.return_value = client_instance
        yield client_instance

@pytest.fixture
def mqtt_service(mock_mqtt_client):
    config = {
        'mqtt': {
            'broker_address': 'localhost',
            'port': 1883,
            'keepalive': 60,
            'username': 'test',
            'password': 'test',
            'debug': False
        }
    }
    service = MQTTService(config)
    return service

def test_init_and_connect(mock_mqtt_client):
    # Arrange
    config = {
        'mqtt': {
            'broker_address': 'test.mosquitto.org',
            'port': 1883,
            'keepalive': 60,
            'username': 'test',
            'password': 'test',
            'debug': False
        }
    }

    # Act
    service = MQTTService(config)

    # Assert
    mock_mqtt_client.username_pw_set.assert_called_once_with('test', 'test')
    mock_mqtt_client.connect.assert_called_once_with('test.mosquitto.org', 1883, 60)
    mock_mqtt_client.loop_start.assert_called_once()

def test_register_device(mqtt_service):
    # Arrange
    device_id = 1
    num_slots = 4
    callback = lambda x: x

    # Act
    mqtt_service.register_device(device_id, num_slots, callback)

    # Assert
    assert mqtt_service.devices[device_id] == num_slots
    assert len(mqtt_service.insertion_callbacks[device_id]) == num_slots
    for slot in range(num_slots):
        assert mqtt_service.insertion_callbacks[device_id][slot] == callback

def test_on_message_callback(mqtt_service):
    # Arrange
    device_id = 1
    slot_id = 0
    cell_id = 123
    callback_mock = MagicMock()
    mqtt_service.register_device(device_id, 1, callback_mock)
    
    msg = MagicMock()
    msg.topic = f"cells_inserted/device_{device_id}"
    msg.payload = json.dumps({
        "slot_id": slot_id,
        "cell_id": cell_id
    }).encode()

    # Act
    mqtt_service.on_message(None, None, msg)

    # Assert
    callback_mock.assert_called_once()
    args = callback_mock.call_args[0][0]
    assert isinstance(args, InsertedCell)
    assert args.device_id == device_id
    assert args.slot_id == slot_id
    assert args.cell_id == cell_id

def test_cell_finished_publishing(mqtt_service, mock_mqtt_client):
    # Arrange
    device_id = "device1"
    slot_id = 0
    cell_id = 123
    capacity = 2500.0
    status = DeviceStatus.DONE
    
    mqtt_service.register_device(device_id, 1)

    # Act
    mqtt_service.cell_finished(device_id, slot_id, cell_id, capacity, status)

    # Assert
    expected_payload = {
        "device_id": device_id,
        "slot_id": slot_id,
        "cell_id": cell_id,
        "capacity": capacity,
        "status": status.name
    }
    mock_mqtt_client.publish.assert_called_once_with(
        f"measurement_done/{device_id}",
        json.dumps(expected_payload)
    )

def test_cleanup(mqtt_service, mock_mqtt_client):
    # Act
    mqtt_service.cleanup()

    # Assert
    mock_mqtt_client.loop_stop.assert_called_once()
    mock_mqtt_client.disconnect.assert_called_once()

def test_debug_mode(mock_mqtt_client):
    # Arrange
    config = {
        'mqtt': {
            'broker_address': 'localhost',
            'port': 1883,
            'keepalive': 60,
            'username': 'test',
            'password': 'test',
            'debug': True
        }
    }

    # Act
    service = MQTTService(config)

    # Assert
    mock_mqtt_client.connect.assert_not_called()
    mock_mqtt_client.loop_start.assert_not_called()