BitStreams

test.py

@@ -1,30 +1,85 @@
 from sampy.raknet.bitstream import Bitstream
 
-a = Bitstream()
-"""
-a.buffer.append(0x42)
-a.buffer.append(0x13)
-a.buffer.append(0xFF)
-a.buffer.append(0xAC)
-a.buffer.append(0x00)
-a.buffer.append(0x69)
-"""
-"""
-a.write(1)
-a.write(0)
-a.write(1)
-a.write(1)
-a.write(1)
-a.write(1)
-a.write(0)
-a.write(1)
+bitstream = Bitstream(b"\x00" * 4)
+MAGIC_BYTES = b"\x01\xff\x02"
+MAGIC_BITS = [int(bit) for bit in " ".join([format(byte, "08b") for byte in MAGIC_BYTES]) if bit in ("0", "1")]
+fail = 0
 
-a.write(1)
-"""
+def print_pretty(prefix, msg):
+    prefix += " " * (40 - len(prefix))
+    print(prefix, msg)
 
-data = [int(x, 16) for x in "00 00 42 90 0b 61 61 61 61 62 62 62 62 63 63 63 63 66 66 66 66 00".split(" ")]
+def reset_bitstream():
+    bitstream._buffer = bytearray(b"\x00" * len(bitstream._buffer))
+    bitstream._offset = 0
 
-for b in data:
-    a.buffer.append(b)
+def test_write_bit(var_bits):
+    global fail
+    print_pretty("<Test bitstream.write_bit:in>", bitstream.pretty())
+    for bit in var_bits:
+        if not bitstream.write_bit(bit):
+            print_pretty("<Test bitstream.write_bit>", "Failed")
+            fail += 1
+    print_pretty("<Test bitstream.write_bit:out>", bitstream.pretty())
 
-print(a)
+def test_read_bit(var_compare_bits):
+    global fail
+    print_pretty("<Test bitstream.read_bit:in>", bitstream.pretty())
+    for bit in var_compare_bits:
+        ok, bit2 = bitstream.read_bit()
+        if not ok:
+            print_pretty("<Test bitstream.read_bit>", "Failed (ok)")
+            fail += 1
+            continue
+        if bit2 != bit:
+            print_pretty("<Test bitstream.read_bit>", "Failed (match)")
+            fail += 1
+            continue
+    print_pretty("<Test bitstream.read_bit:out>", bitstream.pretty())
+
+def test_write(var_bytes):
+    global fail
+    print_pretty("<Test bitstream.write:in>", bitstream.pretty())
+    if not bitstream.write(var_bytes, len(var_bytes) << 3):
+        print_pretty("<Test bitstream.write>", "Failed")
+        fail += 1
+    print_pretty("<Test bitstream.write:out>", bitstream.pretty())
+
+def test_read(var_compare_bytes):
+    global fail
+    print_pretty("<Test bitstream.read:in>", bitstream.pretty())
+    ok, bytez = bitstream.read(len(var_compare_bytes) << 3)
+    if ok:
+        print(var_compare_bytes, bytez)
+        if bytez != var_compare_bytes:
+            print_pretty("<Test bitstream.read>", "Failed (match)")
+            fail += 1
+    else:
+        print_pretty("<Test bitstream.read>", "Failed (ok)")
+        fail += 1
+    print_pretty("<Test bitstream.read:out>", bitstream.pretty())
+
+def test_rw_bit(offset):
+    bitstream._offset = offset
+    test_write_bit(MAGIC_BITS)
+    bitstream._offset = offset
+    test_read_bit(MAGIC_BITS)
+
+def test_rw_bytes(offset):
+    bitstream._offset = offset
+    test_write(MAGIC_BYTES)
+    bitstream._offset = offset
+    test_read(MAGIC_BYTES)
+
+# Test all edges
+for i in range(9):
+    reset_bitstream()
+    test_rw_bit(i)
+
+reset_bitstream()
+
+for i in range(9):
+    reset_bitstream()
+    test_rw_bytes(i)
+
+print("Failed: %d" % fail)