banban/talkingq-url/tests/test_forced_ota.py

from datetime import datetime

import pytest

from banban.service.device import DeviceService
from banban.service.ota_firmware import OtaFirmwareSelection
from handlers.mqtt_handler import TalkingQMQTTService


@pytest.mark.asyncio
async def test_simulated_device_force_ota_dispatches_once_and_completes(monkeypatch):
    device_id = "TalkingQ_XQSN00001005"
    service = DeviceService()
    mqtt_service = TalkingQMQTTService({"broker": "127.0.0.1"})
    record = None
    ota_commands = []
    published = []

    async def fake_select_latest_for_device(status_row):
        assert status_row["device_type"] == 2
        return OtaFirmwareSelection(
            channel_key="gps",
            version="1.8.5",
            firmware_url="https://example.com/gps-1.8.5.bin",
            firmware_id=15,
            force_update=True,
        )

    async def fake_get_firmware_update_dict(requested_device_id):
        assert requested_device_id == device_id
        return dict(record) if record else None

    async def fake_update_firmware_update(*args, **kwargs):
        nonlocal record
        if args:
            kwargs = {"device_id": args[0], **kwargs}
        record = dict(record or {})
        record.update(kwargs)
        record.setdefault("device_id", device_id)
        record.setdefault("updated_at", datetime.now())
        return True

    async def fake_claim_forced_update(**kwargs):
        nonlocal record
        if record and record.get("update_status") in {"sent", "accepted", "updating"}:
            return False
        record = dict(record or {})
        record.update(
            {
                "device_id": device_id,
                "firmware_version": kwargs["target_version"],
                "target_version": kwargs["target_version"],
                "firmware_url": kwargs["firmware_url"],
                "firmware_id": kwargs["firmware_id"],
                "ota_channel": kwargs["ota_channel"],
                "update_status": "sent",
                "progress": 0.0,
                "source": "forced",
                "updated_at": datetime.now(),
            }
        )
        return True

    class FakeMqttService:
        async def send_ota_command(self, requested_device_id, url, version):
            ota_commands.append({"device_id": requested_device_id, "url": url, "version": version})
            return "003"

    async def fake_get_instance():
        return FakeMqttService()

    async def fake_publish(topic, payload):
        published.append((topic, payload))
        return True

    async def fake_schedule_persistence(requested_device_id, label, coro):
        assert requested_device_id == device_id
        await coro

    async def fake_insert_or_update(**kwargs):
        assert kwargs["device_id"] == device_id

    async def fake_get_setting_by_device_id(requested_device_id):
        return {"device_id": requested_device_id, "device_type": 2, "ota_channel": None}

    monkeypatch.setattr(
        "banban.service.device.ota_firmware_service.select_latest_for_device",
        fake_select_latest_for_device,
    )
    monkeypatch.setattr(
        "banban.service.device.device_firmware_update_manager.get_firmware_update_dict",
        fake_get_firmware_update_dict,
    )
    monkeypatch.setattr(
        "banban.service.device.device_firmware_update_manager.update_firmware_update",
        fake_update_firmware_update,
    )
    monkeypatch.setattr(
        "banban.service.device.device_firmware_update_manager.claim_forced_update",
        fake_claim_forced_update,
    )
    monkeypatch.setattr("handlers.mqtt_handler.TalkingQMQTTService.get_instance", fake_get_instance)
    monkeypatch.setattr(mqtt_service, "_publish", fake_publish)
    monkeypatch.setattr(mqtt_service, "_schedule_persistence", fake_schedule_persistence)
    monkeypatch.setattr("handlers.mqtt_handler.device_setting_service.insert_or_update", fake_insert_or_update)
    monkeypatch.setattr("handlers.mqtt_handler.device_setting_service.get_setting_by_device_id", fake_get_setting_by_device_id)
    monkeypatch.setattr(
        "handlers.mqtt_handler.wechat_mp_notification_service.schedule_low_battery_notification",
        lambda **kwargs: None,
    )

    simulated_device_info = {
        "msg_id": "000",
        "data": {
            "id": device_id,
            "version": "1.8.2",
            "device_type": "2",
            "power": 80,
            "signal": 4,
        },
    }
    await mqtt_service._handle_device_info(device_id, simulated_device_info)
    await mqtt_service._handle_device_info(device_id, simulated_device_info)

    assert ota_commands == [{
        "device_id": device_id,
        "url": "https://example.com/gps-1.8.5.bin",
        "version": "1.8.5",
    }]
    assert record["update_status"] == "sent"
    assert record["target_version"] == "1.8.5"
    assert record["source"] == "forced"
    assert published == [
        (f"device/{device_id}/event_resp", {"msg_id": "000", "status": "success"}),
        (f"device/{device_id}/event_resp", {"msg_id": "000", "status": "success"}),
    ]

    await mqtt_service._handle_ota_response(
        device_id,
        {
            "status": "accepted",
            "data": {"current_version": "1.8.2", "target_version": "1.8.5"},
        },
    )
    assert record["update_status"] == "accepted"

    await mqtt_service._handle_ota_response(
        device_id,
        {"status": "updating", "data": {"version": "1.8.5", "progress": 65}},
    )
    assert record["update_status"] == "updating"
    assert record["progress"] == 65.0

    await mqtt_service._handle_ota_response(
        device_id,
        {"status": "success", "data": {"new_version": "1.8.5"}},
    )
    assert record["update_status"] == "success"
    assert record["progress"] == 100.0