bal-electrum-plugin/tests/test_core_qr_transfer.py

"""
Tests for ``bal.core.qrtransfer``.

Covers the BALQR frame encoding used for will transfer via QR codes /
audio modem: encoding, framing, reassembly, malformed input and the preset
list (optionally cross-checked against the ``qrcode`` library's EC-M
capacity when it is installed).

Run:
    source electrum/env/bin/activate
    python3 tests/test_core_qr_transfer.py
"""

import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), os.pardir))

from bal.core.qrtransfer import (
    CHUNK_PRESETS,
    MIN_CHUNK_SIZE,
    InconsistentTotalError,
    MissingFramesError,
    QrTransferError,
    assemble,
    decode_transfer,
    encode_transfer,
    parse_frame,
    preset_index_for_chunk_size,
    split_frames,
)


def _frames(tx_strings, chunk_size, compress=False):
    """Split a payload and return (payload, total, {index: payload})."""
    payload = encode_transfer(tx_strings, compress=compress)
    parsed = {}
    total = None
    for frame in split_frames(payload, chunk_size, compressed=compress):
        t, index, _compressed, p = parse_frame(frame)
        if total is not None:
            assert total == t
        total = t
        parsed[index] = p
    assert total is not None
    return payload, total, parsed


# --------------------------------------------------------------------------- #
# Round trips
# --------------------------------------------------------------------------- #


def test_encode_decode_plain():
    tx_strings = ["00" * 32, "aa" * 40, "ff" * 50]
    payload = encode_transfer(tx_strings, compress=False)
    assert decode_transfer(payload, compressed=False) == tx_strings


def test_encode_decode_compressed():
    tx_strings = ["00" * 32, "aa" * 40, "ff" * 50]
    payload = encode_transfer(tx_strings, compress=True)
    assert decode_transfer(payload, compressed=True) == tx_strings


def test_empty_list_roundtrip():
    assert decode_transfer(encode_transfer([]), compressed=False) == []


# --------------------------------------------------------------------------- #
# Framing
# --------------------------------------------------------------------------- #


def test_single_frame():
    tx_strings = ["11" * 10]
    payload = encode_transfer(tx_strings)
    frames = split_frames(payload, 150)
    assert len(frames) == 1
    total, index, compressed, p = parse_frame(frames[0])
    assert (total, index, compressed) == (1, 1, False)
    assert p == payload


def test_multiple_frames_reassemble():
    tx_strings = ["ab" * 100, "cd" * 100]  # 600 chars -> multiple frames
    _payload, total, parsed = _frames(tx_strings, CHUNK_PRESETS[0][1])
    assert total > 1
    decoded = decode_transfer(assemble(parsed, total), compressed=False)
    assert decoded == tx_strings


def test_size_greater_than_payload():
    tx_strings = ["12" * 5]
    payload = encode_transfer(tx_strings)
    frames = split_frames(payload, 1800)
    assert len(frames) == 1
    _t, _i, _c, p = parse_frame(frames[0])
    assert p == payload


def test_exact_single_frame_boundary():
    # A 138-byte payload exactly fills the 150-byte preset budget (the 12-char
    # empty-flags header plus payload), so the encoded frame is exactly 150.
    tx_strings = ["a" * 138]
    payload = encode_transfer(tx_strings)
    frames = split_frames(payload, 150)
    assert len(frames) == 1
    assert len(frames[0]) == 150
    _t, _i, _c, p = parse_frame(frames[0])
    assert p == payload


def test_frames_fit_chunk_size():
    tx_strings = ["".join("{:02x}".format(i) * 2) for i in range(300)]
    payload = encode_transfer(tx_strings)
    for _label, size in CHUNK_PRESETS:
        for frame in split_frames(payload, size):
            assert len(frame) <= size, (size, len(frame))


def test_single_tx_larger_than_chunk():
    # A huge serialized tx must be split over several frames and reassemble
    # exactly (positional slicing is safe for hex/base64 text).
    tx_strings = ["7b" * 1000]  # 2000 chars
    payload = encode_transfer(tx_strings)
    frames = split_frames(payload, 150)
    assert len(frames) > 1
    parsed = {parse_frame(f)[1]: parse_frame(f)[3] for f in frames}
    total = parse_frame(frames[0])[0]
    assert assemble(parsed, total) == payload


def test_compressed_frames_carry_flag():
    tx_strings = ["ab" * 40]
    frames = split_frames(encode_transfer(tx_strings, compress=True), 150, compressed=True)
    for frame in frames:
        _t, _i, compressed, _p = parse_frame(frame)
        assert compressed is True
    # Plain frames do not.
    frames_plain = split_frames(encode_transfer(tx_strings), 150)
    _t, _i, compressed, _p = parse_frame(frames_plain[0])
    assert compressed is False


def test_compressed_roundtrip_through_frames():
    tx_strings = ["ab" * 50, "cd" * 50, "12" * 60]
    payload = encode_transfer(tx_strings, compress=True)
    parsed = {}
    total = None
    for frame in split_frames(payload, 400, compressed=True):
        t, index, _c, p = parse_frame(frame)
        total = t
        parsed[index] = p
    assert total is not None
    decoded = decode_transfer(assemble(parsed, total), compressed=True)
    assert decoded == tx_strings


# --------------------------------------------------------------------------- #
# Malformed input
# --------------------------------------------------------------------------- #


def test_parse_bad_magic_and_version():
    for frame in (
        "BALQR|1|1||a",  # missing version
        "BALQR2|1|1||a",  # unknown version
        "XXXXX1|1|1||a",  # unknown magic
    ):
        try:
            parse_frame(frame)
        except QrTransferError:
            pass
        else:
            raise AssertionError("expected QrTransferError for: {}".format(frame))


def test_parse_bad_arity():
    for frame in ("BALQR1", "BALQR1|1|1|"):
        try:
            parse_frame(frame)
        except QrTransferError:
            pass
        else:
            raise AssertionError("expected QrTransferError for: {}".format(frame))


def test_parse_pipe_in_payload_is_folded():
    # maxsplit keeps the tail (including any inner '|') in the payload part.
    frame = "BALQR1|1|1||a|b|c"
    total, index, compressed, payload = parse_frame(frame)
    assert (total, index, compressed) == (1, 1, False)
    assert payload == "a|b|c"


def test_parse_bad_numbers():
    for frame in (
        "BALQR1|x|1||a",
        "BALQR1|1|y||a",
        "BALQR1|0|1||a",
        "BALQR1|1|0||a",
        "BALQR1|1|2||a",  # index beyond total
        "BALQR1|-1|1||a",
    ):
        try:
            parse_frame(frame)
        except QrTransferError:
            pass
        else:
            raise AssertionError("expected QrTransferError for: {}".format(frame))


def test_parse_bad_flags():
    try:
        parse_frame("BALQR1|1|1|Q|payload")
    except QrTransferError:
        pass
    else:
        raise AssertionError("expected QrTransferError for unknown flags")


def test_assemble_missing_frames():
    try:
        assemble({1: "a", 3: "c"}, total=3)
    except MissingFramesError as e:
        assert e.missing == [2]
    else:
        raise AssertionError("expected MissingFramesError")


def test_assemble_index_beyond_total():
    try:
        assemble({1: "a", 2: "b"}, total=1)
    except InconsistentTotalError:
        pass
    else:
        raise AssertionError("expected InconsistentTotalError")


def test_assemble_order_and_total_validation():
    assert assemble({1: "a", 2: "b"}, total=2) == "ab"
    try:
        assemble({}, total=0)
    except QrTransferError:
        pass
    else:
        raise AssertionError("expected QrTransferError")


# --------------------------------------------------------------------------- #
# Constants / presets
# --------------------------------------------------------------------------- #


def test_preset_count_and_order():
    assert len(CHUNK_PRESETS) == 4
    budgets = [budget for _label, budget in CHUNK_PRESETS]
    assert budgets == sorted(budgets)


def test_preset_index_for_chunk_size():
    for index, (_label, budget) in enumerate(CHUNK_PRESETS):
        assert preset_index_for_chunk_size(budget) == index
    assert preset_index_for_chunk_size(150) == 0
    assert preset_index_for_chunk_size(1800) == 3


def test_min_chunk_size_guard():
    try:
        split_frames("x" * 10, MIN_CHUNK_SIZE - 1)
    except QrTransferError:
        pass
    else:
        raise AssertionError("expected QrTransferError for tiny chunk size")


def test_split_frame_headers_consistent():
    tx_strings = ["ab" * 80]
    payload = encode_transfer(tx_strings)
    frames = split_frames(payload, 150)
    totals = {parse_frame(frame)[0] for frame in frames}
    assert len(totals) == 1
    assert totals.pop() == len(frames)


# --------------------------------------------------------------------------- #
# Optional: cross-check presets against the qrcode library (EC level M)
# --------------------------------------------------------------------------- #


def test_presets_fit_qrcode_ec_m():
    """Every preset budget must render inside a QR at EC level M."""
    try:
        import qrcode

        from qrcode.constants import ERROR_CORRECT_M
    except ImportError:
        print("qrcode not installed - skipping capacity check")
        return
    for _label, size in CHUNK_PRESETS:
        # Worst-case frame: header with the largest plausible total/index plus
        # a full payload of the preset budget.
        frame = "BALQR1|9999|9999|Z|" + "a" * (size - 14)
        qr = qrcode.QRCode(error_correction=ERROR_CORRECT_M, border=2)
        qr.add_data(frame)
        qr.get_matrix()  # raises DataOverflowError if it does not fit


# --------------------------------------------------------------------------- #
if __name__ == "__main__":
    import traceback

    failures = 0
    for _name, fn in sorted(globals().items()):
        if _name.startswith("test_") and callable(fn):
            try:
                fn()
                print("ok: {}".format(_name))
            except Exception:
                failures += 1
                print("FAIL: {}".format(_name))
                traceback.print_exc()
    if failures:
        print("{} test(s) failed".format(failures))
        sys.exit(1)
    print("all tests passed")