plm-lottery/tests/unit/test_electrum_client.py

import asyncio
import json

from app.electrum.client import ElectrumClient, ElectrumError


class FakeWriter:
    def __init__(self):
        self.written = b""

    def write(self, data: bytes) -> None:
        self.written += data

    async def drain(self) -> None:
        pass

    def close(self) -> None:
        pass

    async def wait_closed(self) -> None:
        pass


async def _client_with_fake_transport() -> tuple[ElectrumClient, asyncio.StreamReader, FakeWriter]:
    client = ElectrumClient("localhost", 1234)
    reader = asyncio.StreamReader()
    writer = FakeWriter()
    client._reader = reader
    client._writer = writer
    client._read_task = asyncio.create_task(client._read_loop())
    return client, reader, writer


async def test_request_resolves_on_matching_response():
    client, reader, writer = await _client_with_fake_transport()

    task = asyncio.create_task(client.request("blockchain.headers.subscribe"))
    await asyncio.sleep(0)  # let request() write the payload
    sent = json.loads(writer.written.decode())
    assert sent["method"] == "blockchain.headers.subscribe"

    reader.feed_data((json.dumps({"id": sent["id"], "result": {"height": 100}}) + "\n").encode())
    result = await task
    assert result == {"height": 100}

    client._read_task.cancel()


async def test_error_response_raises_electrum_error():
    client, reader, writer = await _client_with_fake_transport()

    task = asyncio.create_task(client.request("blockchain.transaction.broadcast", ["deadbeef"]))
    await asyncio.sleep(0)
    sent = json.loads(writer.written.decode())

    reader.feed_data((json.dumps({"id": sent["id"], "error": "bad tx"}) + "\n").encode())
    try:
        await task
        assert False, "expected ElectrumError"
    except ElectrumError:
        pass

    client._read_task.cancel()


async def test_notification_delivered_to_subscription_queue():
    client, reader, writer = await _client_with_fake_transport()
    queue = client.notifications("blockchain.scripthash.subscribe")

    push = {"method": "blockchain.scripthash.subscribe", "params": ["abcd", "newstatus"]}
    reader.feed_data((json.dumps(push) + "\n").encode())
    await asyncio.sleep(0)

    params = await asyncio.wait_for(queue.get(), timeout=1)
    assert params == ["abcd", "newstatus"]

    client._read_task.cancel()


async def test_request_times_out_instead_of_hanging_forever(monkeypatch):
    """B-01: a server that owes us a reply and never sends one used to hang the
    caller permanently — which meant a POST /bets could hang while holding the
    per-user lock, and the confirmation poller could stop polling for good."""
    from app.electrum import client as client_module

    monkeypatch.setattr(client_module, "_REQUEST_TIMEOUT_SECONDS", 0.05)
    client, reader, writer = await _client_with_fake_transport()

    try:
        await client.request("blockchain.transaction.get", ["deadbeef"])
        assert False, "expected ElectrumError"
    except ElectrumError as exc:
        assert "timed out" in str(exc)

    # The connection is torn down, so callers stop reusing a server that owes us.
    assert client._closed.is_set()
    client._read_task.cancel()


async def test_wait_closed_resolves_when_the_read_loop_dies():
    """B-01: the read loop dying used to be invisible — the listener sat on its
    notification queues forever and never reconnected."""
    client, reader, writer = await _client_with_fake_transport()

    waiter = asyncio.create_task(client.wait_closed())
    reader.feed_eof()  # peer closed the connection

    await asyncio.wait_for(waiter, timeout=1)
    client._read_task.cancel()


async def test_pending_request_fails_when_the_connection_drops():
    client, reader, writer = await _client_with_fake_transport()

    task = asyncio.create_task(client.request("server.ping"))
    await asyncio.sleep(0)
    reader.feed_eof()

    try:
        await asyncio.wait_for(task, timeout=1)
        assert False, "expected ElectrumError"
    except ElectrumError:
        pass


def test_parse_endpoints_puts_the_primary_first_and_dedupes():
    from app.electrum.client import ElectrumEndpoint, parse_endpoints

    endpoints = parse_endpoints(
        "primary.example", 50002, True, "second.example:50002, third.example:50001:notls, primary.example:50002"
    )
    assert endpoints == [
        ElectrumEndpoint("primary.example", 50002, True),
        ElectrumEndpoint("second.example", 50002, True),
        ElectrumEndpoint("third.example", 50001, False),
    ]


def test_parse_endpoints_handles_an_empty_fallback_list():
    from app.electrum.client import ElectrumEndpoint, parse_endpoints

    assert parse_endpoints("only.example", 50002, True, "") == [ElectrumEndpoint("only.example", 50002, True)]
    assert parse_endpoints("only.example", 50002, True, "  ,  ") == [
        ElectrumEndpoint("only.example", 50002, True)
    ]


def test_parse_endpoints_rejects_a_malformed_entry():
    """A typo in a fallback server must fail at startup, not during the outage when
    the fallback is the thing that's needed."""
    import pytest

    from app.electrum.client import parse_endpoints

    for bad in ["nohost:", "host:notaport", "host", "host:50002:weird"]:
        with pytest.raises(ValueError):
            parse_endpoints("primary.example", 50002, True, bad)


# --- B-75 (ELC-20): an oversized or malformed reply is a request error, never a
# dead connection. A tiny line limit stands in for 16 MiB. ---------------------


async def _send_request(client, writer, reader):
    task = asyncio.create_task(client.request("blockchain.scripthash.get_history", ["ab"]))
    await asyncio.sleep(0)
    return task, json.loads(writer.written.decode().splitlines()[-1])["id"]


async def test_the_line_limit_is_16_mib_and_connect_applies_it(monkeypatch):
    """ELC-20 / B-75: the StreamReader is opened with a 16 MiB limit, not 64 KiB."""
    from app.electrum import client as client_module

    assert client_module._MAX_LINE_BYTES == 16 * 1024 * 1024
    seen = {}

    async def fake_open_connection(host, port, **kwargs):
        seen.update(kwargs)
        raise OSError("stop here")

    monkeypatch.setattr(asyncio, "open_connection", fake_open_connection)
    try:
        await ElectrumClient("localhost", 1).connect()
    except OSError:
        pass
    assert seen["limit"] == 16 * 1024 * 1024


async def test_an_oversized_reply_fails_its_request_and_keeps_the_connection(monkeypatch):
    """ELC-20 / B-75: the reply is discarded to its newline (even when it arrives in
    pieces), the request whose id is in its head fails at once, and the next reply
    on the same connection is still delivered."""
    from app.electrum import client as client_module

    monkeypatch.setattr(client_module, "_MAX_LINE_BYTES", 1024)
    client, reader, writer = await _client_with_fake_transport()

    task, request_id = await _send_request(client, writer, reader)
    head = b'{"id": %d, "result": [' % request_id
    reader.feed_data(head + b"1," * 600)  # already over the limit, no newline yet
    await asyncio.sleep(0)
    reader.feed_data(b"1," * 600 + b"1]}\n")
    try:
        await asyncio.wait_for(task, timeout=1)
        assert False, "expected ElectrumError"
    except ElectrumError:
        pass
    assert not client._closed.is_set()

    task2, request_id2 = await _send_request(client, writer, reader)
    reader.feed_data((json.dumps({"id": request_id2, "result": "ok"}) + "\n").encode())
    assert await asyncio.wait_for(task2, timeout=1) == "ok"

    client._read_task.cancel()


async def test_an_oversized_reply_with_the_id_at_the_end_also_fails_its_request(monkeypatch):
    """ELC-20: the id is looked for in the last bytes of the line too."""
    from app.electrum import client as client_module

    monkeypatch.setattr(client_module, "_MAX_LINE_BYTES", 1024)
    client, reader, writer = await _client_with_fake_transport()

    task, request_id = await _send_request(client, writer, reader)
    reader.feed_data(b'{"result": [' + b"1," * 1000 + b'1], "id": %d}\n' % request_id)
    try:
        await asyncio.wait_for(task, timeout=1)
        assert False, "expected ElectrumError"
    except ElectrumError:
        pass
    assert not client._closed.is_set()

    client._read_task.cancel()


async def test_a_reply_that_is_not_json_fails_its_request_not_the_connection():
    """ELC-20 / B-75: json.loads used to raise out of the read loop and close it."""
    client, reader, writer = await _client_with_fake_transport()

    task, request_id = await _send_request(client, writer, reader)
    reader.feed_data(b'{"id": %d, "result": [1, 2,\n' % request_id)
    try:
        await asyncio.wait_for(task, timeout=1)
        assert False, "expected ElectrumError"
    except ElectrumError:
        pass
    assert not client._closed.is_set()

    # Garbage with no readable id, and a non-object message, are dropped silently.
    reader.feed_data(b"not json at all\n[1, 2]\n")
    task2, request_id2 = await _send_request(client, writer, reader)
    reader.feed_data((json.dumps({"id": request_id2, "result": 7}) + "\n").encode())
    assert await asyncio.wait_for(task2, timeout=1) == 7

    client._read_task.cancel()