d8208685a2
Reactor/loop-pool engine with TCP/µTP/MSE transports, per-connection pipelining, priority-driven piece selection with endgame, and the Python FFI test harness. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
89 lines
2.7 KiB
Python
89 lines
2.7 KiB
Python
"""
|
|
Regression tests for swarm_download's endgame helpers.
|
|
|
|
Run with: python tests/test_swarm_endgame.py
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import os
|
|
import sys
|
|
|
|
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, os.path.join(ROOT, "harness"))
|
|
|
|
from engine_ffi import BLOCK_SIZE # noqa: E402
|
|
from swarm_download import EndgameController, PieceAssembler # noqa: E402
|
|
|
|
|
|
class FakeMeta:
|
|
def __init__(self, payloads: list[bytes]):
|
|
self.payloads = payloads
|
|
self.num_pieces = len(payloads)
|
|
self.piece_length = max(len(p) for p in payloads)
|
|
self.total_size = sum(len(p) for p in payloads)
|
|
self.piece_hashes = [hashlib.sha1(p).digest() for p in payloads]
|
|
|
|
def piece_len(self, piece: int) -> int:
|
|
return len(self.payloads[piece])
|
|
|
|
|
|
class FakeEngine:
|
|
def __init__(self):
|
|
self.priorities = []
|
|
self.rearms = []
|
|
|
|
def set_priority(self, tid: int, piece: int, priority: int) -> None:
|
|
self.priorities.append((tid, piece, priority))
|
|
|
|
def request_piece(self, tid: int, piece: int) -> None:
|
|
self.rearms.append((tid, piece))
|
|
|
|
|
|
def test_piece_assembler_ignores_duplicate_blocks():
|
|
payload = (b"a" * BLOCK_SIZE) + b"tail"
|
|
meta = FakeMeta([payload])
|
|
done = bytearray(meta.num_pieces)
|
|
asm = PieceAssembler(meta, done)
|
|
|
|
first = payload[:BLOCK_SIZE]
|
|
tail = payload[BLOCK_SIZE:]
|
|
|
|
added, complete = asm.add_block(0, 0, first)
|
|
assert added
|
|
assert not complete
|
|
assert asm.received[0] == len(first)
|
|
|
|
added, complete = asm.add_block(0, 0, first)
|
|
assert not added
|
|
assert not complete
|
|
assert asm.received[0] == len(first)
|
|
|
|
added, complete = asm.add_block(0, BLOCK_SIZE, tail)
|
|
assert added
|
|
assert complete
|
|
assert hashlib.sha1(asm.piece_bytes(0)).digest() == meta.piece_hashes[0]
|
|
|
|
|
|
def test_endgame_rearms_only_unfinished_pieces_on_interval():
|
|
meta = FakeMeta([b"a", b"b", b"c", b"d"])
|
|
done = bytearray([0, 1, 0, 0])
|
|
eng = FakeEngine()
|
|
ctl = EndgameController(meta, min_pieces=3, peer_factor=2.0, interval=3.0)
|
|
|
|
ctl.maybe_rearm(eng, 7, done, done_count=1, connected=2, now=10.0)
|
|
assert eng.priorities == [(7, 0, 255), (7, 2, 255), (7, 3, 255)]
|
|
assert eng.rearms == [(7, 0), (7, 2), (7, 3)]
|
|
|
|
ctl.maybe_rearm(eng, 7, done, done_count=1, connected=2, now=11.0)
|
|
assert eng.rearms == [(7, 0), (7, 2), (7, 3)]
|
|
|
|
done[2] = 1
|
|
ctl.maybe_rearm(eng, 7, done, done_count=2, connected=2, now=13.0)
|
|
assert eng.rearms[-2:] == [(7, 0), (7, 3)]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_piece_assembler_ignores_duplicate_blocks()
|
|
test_endgame_rearms_only_unfinished_pieces_on_interval()
|
|
print("swarm endgame OK")
|