ohbf/tests/test_ohbf.py

import math
import random
import warnings

import pytest

from ohbf import (
    OneHashBloomFilter,
    bloom_parameters,
    choose_partitions,
    partition_positions,
    primes_up_to,
)

HAND_HASHES = {b"apple": 100, b"pear": 37, b"plum": 58, b"grape": 345}


def hand_filter() -> OneHashBloomFilter:
    return OneHashBloomFilter((5, 7, 11), hash_fn=HAND_HASHES.__getitem__)


def test_no_false_negatives():
    f = OneHashBloomFilter.from_capacity(5_000, 0.01)
    keys = [f"key-{i}" for i in range(5_000)]
    for key in keys:
        f.add(key)
    assert all(key in f for key in keys)
    assert len(f) == 5_000


def test_str_and_bytes_keys_are_equivalent():
    f = OneHashBloomFilter.from_capacity(100, 0.01)
    f.add("héllo")
    assert "héllo".encode() in f
    assert f.positions("héllo") == f.positions("héllo".encode())


def test_rejects_non_coprime_sizes():
    with pytest.raises(ValueError, match="coprime"):
        OneHashBloomFilter((6, 9, 11))


@pytest.mark.parametrize("sizes", [(), (1, 5), (5, 7.0)])
def test_rejects_bad_sizes(sizes):
    with pytest.raises(ValueError):
        OneHashBloomFilter(sizes)


def test_rejects_other_key_types():
    with pytest.raises(TypeError):
        OneHashBloomFilter((5, 7)).add(42)


def test_primes_up_to():
    assert primes_up_to(1) == []
    assert primes_up_to(30) == [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]


@pytest.mark.parametrize(
    "m_target,k", [(10, 3), (1_000, 4), (95_851, 7), (958_506, 7), (10**6, 32)]
)
def test_choose_partitions_distinct_primes_near_target(m_target, k):
    parts = choose_partitions(m_target, k)
    assert len(parts) == k
    assert len(set(parts)) == k
    primes = set(primes_up_to(max(parts)))
    assert all(p in primes for p in parts)
    assert parts == sorted(parts)
    # consecutive: no prime is skipped inside the window
    assert parts == [p for p in sorted(primes) if parts[0] <= p <= parts[-1]]
    assert abs(sum(parts) - m_target) / m_target < 0.02 or m_target < 100


def test_choose_partitions_reproduces_paper_table_ii():
    table_ii = {
        10_000: [971, 977, 983, 991, 997, 1009, 1013, 1019, 1021, 1031],
        20_000: [1973, 1979, 1987, 1993, 1997, 1999, 2003, 2011, 2017, 2027],
        40_000: [3947, 3967, 3989, 4001, 4003, 4007, 4013, 4019, 4021, 4027],
        80_000: [7949, 7951, 7963, 7993, 8009, 8011, 8017, 8039, 8053, 8059],
        160_000: [15937, 15959, 15971, 15973, 15991, 16001, 16007, 16033, 16057, 16061],
        320_000: [31957, 31963, 31973, 31981, 31991, 32003, 32009, 32027, 32029, 32051],
    }
    for m_target, expected in table_ii.items():
        assert choose_partitions(m_target, 10) == expected


def test_choose_partitions_is_closest_window():
    for m_target, k in [(500, 3), (12_345, 5), (77_777, 8)]:
        parts = choose_partitions(m_target, k)
        primes = primes_up_to(2 * m_target)
        sums = [sum(primes[i : i + k]) for i in range(len(primes) - k + 1)]
        assert abs(sum(parts) - m_target) == min(abs(s - m_target) for s in sums)


def test_hand_example_positions():
    f = hand_filter()
    assert f.positions("apple") == (0, 2, 1)
    assert f.positions("pear") == (2, 2, 4)
    assert f.positions("grape") == (0, 2, 4)


def test_hand_example_membership_and_false_positive():
    f = hand_filter()
    f.add("apple")
    f.add("pear")
    assert "apple" in f
    assert "plum" not in f
    assert "grape" in f  # never added: false positive
    assert f.partition_bits(0) == [1, 0, 1, 0, 0]
    assert f.partition_bits(1) == [0, 0, 1, 0, 0, 0, 0]
    assert f.partition_bits(2) == [0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0]
    assert f.set_bit_count() == 5


def test_empirical_fp_rate_matches_eq6():
    rng = random.Random(2015)
    f = OneHashBloomFilter.from_capacity(10_000, 0.01)
    for _ in range(10_000):
        f.add(b"in:" + rng.randbytes(16))
    queries = 100_000
    hits = sum(b"out:" + rng.randbytes(16) in f for _ in range(queries))
    observed = hits / queries
    expected = f.expected_fp_rate()
    assert expected == pytest.approx(0.01, rel=0.05)
    # binomial std at p=0.01, N=1e5 is ~3.1e-4; allow ~5 sigma
    sigma = math.sqrt(expected * (1 - expected) / queries)
    assert abs(observed - expected) < 5 * sigma


def test_coprimality_matters():
    rng = random.Random(7)
    hashes = [rng.getrandbits(128) for _ in range(5_000)]

    # 5 | 10, so h % 10 fixes h % 5: only 10 of the 50 pairs ever occur.
    bad = {partition_positions(h, (5, 10)) for h in hashes}
    assert len(bad) < 5 * 10
    assert len(bad) == 10
    assert all(a == b % 5 for a, b in bad)

    good = {partition_positions(h, (5, 7)) for h in hashes}
    assert len(good) == 5 * 7


def test_digest_size_covers_partition_product():
    f = OneHashBloomFilter((136883, 136889, 136897))
    needed = math.prod(f.sizes).bit_length() + 64
    assert 8 * f.digest_size >= needed
    assert f.positions("x") == f.positions("x")  # deterministic


def test_warns_when_512_bits_is_not_enough():
    sizes = choose_partitions(10**6, 32)  # product is ~32 * 15 = 480+ bits
    assert math.prod(sizes).bit_length() + 64 > 512
    with pytest.warns(UserWarning, match="independence"):
        f = OneHashBloomFilter(sizes)
    assert f.digest_size == 64
    f.add("still works")
    assert "still works" in f


def test_no_warning_with_injected_hash():
    with warnings.catch_warnings():
        warnings.simplefilter("error")
        OneHashBloomFilter(choose_partitions(10**6, 32), hash_fn=lambda b: 0)


def test_bloom_parameters():
    m, k = bloom_parameters(10_000, 0.01)
    assert m == 95_851
    assert k == 7
    with pytest.raises(ValueError):
        bloom_parameters(0, 0.01)
    with pytest.raises(ValueError):
        bloom_parameters(10, 1.5)


def test_expected_fp_rate_empty_and_explicit_n():
    f = OneHashBloomFilter((5, 7, 11))
    assert f.expected_fp_rate() == 0.0
    assert f.expected_fp_rate(n=2) == pytest.approx(
        (1 - (4 / 5) ** 2) * (1 - (6 / 7) ** 2) * (1 - (10 / 11) ** 2)
    )