KX-Bridge-Release/tests/test_client_robustness.py

"""Robustness fixes found during a targeted code review of kobrax_client.py:

1. _drain()/_dispatch() must not get stuck reprocessing the same malformed
   packet forever when the printer sends valid JSON that isn't an object
   (e.g. a bare number or list) - previously an exception from _dispatch()
   escaped before self._buf was advanced, so the same bytes sat at the
   front of the buffer and crashed every subsequent _drain() call, each one
   forcing a reconnect via the reader loop's exception handler.
2. publish()'s pending-dict registration/cleanup must be safe against
   concurrent callers for the same msg_type, and a stale/late reply must
   not be delivered to a newer, unrelated caller waiting on the same
   report-topic suffix.
3. upload_gcode() must raise a clear error for a malformed upload URL
   (missing "?s=") instead of silently sending an unauthenticated request,
   and must not leak the upload socket on a send/recv failure.
"""
import argparse
import json
import socket
import tempfile
import threading
import time
from unittest.mock import MagicMock, patch

import pytest

from kobrax_client import KobraXClient, _build_publish


def _client(**overrides):
    kwargs = dict(
        host="192.168.1.100", username="u", password="p",
        mode_id="20030", device_id="abc123", port=9883,
        client_id="test",
    )
    kwargs.update(overrides)
    return KobraXClient(**kwargs)


def _frame(topic: str, payload) -> bytes:
    """Builds a raw MQTT PUBLISH frame carrying `payload` as JSON, the same
    wire format _drain() parses."""
    return _build_publish(topic, json.dumps(payload))


# --- Fix 1: malformed (non-dict) JSON payload must not wedge the buffer ---

def test_drain_advances_buffer_past_non_dict_payload():
    c = _client()
    # A bare JSON number is valid JSON but not a dict - json.loads succeeds,
    # but _dispatch()'s dict-oriented logic would previously crash on it.
    bad = _frame("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report", 42)
    good = _frame("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report", {"state": "done"})
    c._buf = bad + good

    c._drain()  # must not raise

    # Both packets must have been consumed - the malformed one logged/skipped,
    # the buffer advanced past it so the following valid packet is also
    # processed, not left stuck behind it.
    assert c._buf == b""


def test_drain_does_not_reprocess_bad_packet_on_repeated_calls():
    """Regression guard for the original bug: before the fix, a crash in
    _dispatch() left self._buf untouched, so the bad packet stayed at the
    front and every _drain() call re-crashed on it."""
    c = _client()
    bad = _frame("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report", [1, 2, 3])
    c._buf = bad

    c._drain()
    assert c._buf == b""  # consumed, not stuck

    # A second call on the now-empty buffer must be a no-op, not a re-crash.
    c._drain()
    assert c._buf == b""


def test_dispatch_rejects_non_dict_payload_directly():
    c = _client()
    c._dispatch("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report", "not a dict")
    # No exception, no registered pending entries get corrupted.
    assert c._pending_report == {}


# --- Fix 2: pending-dict thread-safety + msgid correlation ---

def test_publish_concurrent_same_msg_type_both_get_delivered():
    """Two overlapping publish() calls for the same msg_type must each
    receive their OWN reply, not have one silently miss out because the
    other already claimed the shared report_key slot."""
    c = _client()
    c._running = True
    c._sock = MagicMock()
    c._ensure_reader = lambda: None  # no real reader thread needed for this test

    results = {}

    def call(label):
        results[label] = c.publish("info", "query", timeout=2.0)

    t1 = threading.Thread(target=call, args=("a",))
    t2 = threading.Thread(target=call, args=("b",))
    t1.start()
    time.sleep(0.02)
    t2.start()
    time.sleep(0.05)

    # Simulate the printer replying to whichever msgid-bearing requests are
    # currently pending, by msgid (the reliable path both requests can use).
    with c._pending_lock:
        pending = dict(c._pending_msgid)
    for msgid, entry in pending.items():
        c._dispatch("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report",
                    {"msgid": msgid, "state": "done", "data": {"ok": True}})

    t1.join(timeout=3)
    t2.join(timeout=3)

    assert results["a"] is not None
    assert results["b"] is not None


def test_dispatch_ignores_stale_reply_with_mismatched_msgid():
    """A late reply carrying a msgid that doesn't match what the current
    report_key waiter is expecting must not be delivered to it - it belongs
    to an earlier, already-resolved/timed-out request."""
    c = _client()
    c._running = True
    c._sock = MagicMock()

    event = threading.Event()
    entry = {"event": event, "result": None, "msgid": "the-real-one"}
    with c._pending_lock:
        c._pending_report["info/report"] = entry

    # A stale reply for a different msgid must not resolve this waiter.
    c._dispatch("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report",
                {"msgid": "some-other-stale-id", "state": "done"})
    assert not event.is_set()
    assert entry["result"] is None

    # The matching reply must resolve it.
    c._dispatch("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report",
                {"msgid": "the-real-one", "state": "done"})
    assert event.is_set()
    assert entry["result"]["msgid"] == "the-real-one"


def test_dispatch_delivers_reply_without_msgid_as_before():
    """Regression guard: plenty of printer reports carry no msgid at all
    (e.g. spontaneous status pushes) - those must still resolve a waiter
    registered without one (entry["msgid"] is falsy)."""
    c = _client()
    event = threading.Event()
    entry = {"event": event, "result": None, "msgid": None}
    with c._pending_lock:
        c._pending_report["info/report"] = entry

    c._dispatch("anycubic/anycubicCloud/v1/printer/public/20030/abc123/info/report",
                {"state": "done"})
    assert event.is_set()


# --- Fix 3: upload_gcode() URL validation + socket cleanup ---

def test_upload_gcode_raises_clear_error_on_missing_session_token(tmp_path):
    c = _client()
    f = tmp_path / "test.gcode"
    f.write_text("; gcode")

    with pytest.raises(RuntimeError, match="session token"):
        c.upload_gcode(str(f), upload_url="http://192.168.1.100:18910/gcode_upload")


def test_upload_gcode_closes_socket_on_send_failure(tmp_path):
    c = _client()
    f = tmp_path / "test.gcode"
    f.write_text("; gcode")

    fake_sock = MagicMock()
    fake_sock.sendall.side_effect = ConnectionResetError("printer went away")

    with patch("socket.create_connection", return_value=fake_sock):
        with pytest.raises(ConnectionResetError):
            c.upload_gcode(str(f), upload_url="http://192.168.1.100:18910/gcode_upload?s=tok123")

    fake_sock.close.assert_called_once()