Allele Frequency Count

"""Local-loop equivalence test for the `allele_frequency_count` flagship bundle.Proves the full pinned pipeline end-to-end on synthetic contributors:    keygen -> encode -> encrypt (N>=3) -> compute -> decrypt -> decodeand asserts, per docs/simulation_mode.md's oracle claim:    decode(decrypt(compute(encrypt(encode(raw)))))  ==  cleartext aggregate**bit-exact** (BFV, precision_tolerance 0), AND that the append-1 sentineldecrypts to **exactly N** (docs/spec.md).The pure functions live in the author modules (server.py / local_project_owner.py/ local_data_owner.py), grouped by role per docs/rfcs/0002. The numbered stagefiles at the bundle root are kit-owned argparse shims that call these samefunctions — so testing the functions directly exercises the identical logic theshims (and the hosted worker, for compute) run. Run:    uv --project signed/env run --group dev pytest tests/            # from bundle root    # or, with tenseal already importable:    python -m pytest tests/If TenSEAL cannot be imported the whole module skips with a clear reason (realcode, no pseudo-code — the skip is only for a machine that cannot install thesealed dependency)."""from __future__ import annotationsimport jsonimport pathlibimport randomimport sysimport pytestBUNDLE_ROOT = pathlib.Path(__file__).resolve().parent.parent / "signed"# The author modules live at the bundle root (importable names, unlike the# digit-prefixed shim files). Put the bundle root first on sys.path so# `import server` resolves to THIS bundle's server.py.sys.path.insert(0, str(BUNDLE_ROOT))# The encrypted engine needs TenSEAL. Skip cleanly (not fail) if it is absent —# but note this in the task's `unresolved` if it ever fires.pytest.importorskip("tenseal", reason="TenSEAL not installed; sealed env not built")import local_data_owner  # noqa: E402  (after sys.path insert)import local_project_owner  # noqa: E402import server  # noqa: E402# HomomorphicEncryption.org coeff-modulus caps (max Σ coeff_mod_bit_sizes) at N.# achieved(N, Σ) = strictest level L with Σ <= CAP[N][L]. Computed by the harness,# NEVER read back from SEAL (SEAL only validates at tc128)._CAP = {    8192: {256: 118, 192: 152, 128: 218},    16384: {256: 237, 192: 305, 128: 438},    32768: {256: 476, 192: 611, 128: 881},}def _achieved_security(poly_modulus_degree: int, coeff_bits: list[int]) -> int:    """The strictest security level whose cap the coeff chain fits under."""    total = sum(coeff_bits)    caps = _CAP[poly_modulus_degree]    for level in (256, 192, 128):        if total <= caps[level]:            return level    raise AssertionError(f"Σ={total} exceeds even the 128 cap at N={poly_modulus_degree}")def _cleartext_aggregate(raw_vectors: list[list], length: int) -> list[int]:    """The correctness oracle: sum encoded vectors coordinate-wise in cleartext."""    counts = [0] * length    for raw in raw_vectors:        for j, value in enumerate(local_data_owner.encode(raw, length)):            counts[j] += value    return countsdef _run_pipeline(raw_vectors: list[list], length: int, security: int = 128) -> dict:    """keygen -> encode -> encrypt -> compute (server) -> decrypt -> decode."""    secret_ctx, public_ctx = local_project_owner.keygen(security=security)    # Local data-owner stages, once per contributor.    ciphertexts = []    for raw in raw_vectors:        encoded = local_data_owner.encode(raw, length)        ciphertexts.append(local_data_owner.encrypt(public_ctx, encoded))    # The ONLY server-side stage: sum ciphertexts under the PUBLIC context.    result_ct = server.compute(ciphertexts, public_ctx)    # Local project-owner stages: decrypt with the secret context, then decode.    plain = local_project_owner.decrypt(secret_ctx, result_ct)    return local_project_owner.decode(plain, length)def test_local_loop_matches_cleartext_and_fixtures():    """The committed 4-contributor fixtures decode to the committed expected."""    vectors_dir = BUNDLE_ROOT / "tests" / "vectors"    expected = json.loads((BUNDLE_ROOT / "tests" / "expected" / "aggregate.json").read_text())    length = expected["coordinates_length"]    raw_vectors = [        json.loads(path.read_text()) for path in sorted(vectors_dir.glob("*.json"))    ]    assert len(raw_vectors) >= 3, "need >=3 synthetic contributors"    result = _run_pipeline(raw_vectors, length)    # Sentinel recovers the EXACT contributor count.    assert result["n_contributors"] == len(raw_vectors)    assert result["n_contributors"] == expected["n_contributors"]    # Encrypted aggregate == cleartext oracle == committed fixture, bit-exact.    assert result["allele_counts"] == _cleartext_aggregate(raw_vectors, length)    assert result["allele_counts"] == expected["allele_counts"]    assert result["allele_frequencies"] == pytest.approx(        expected["allele_frequencies"], abs=0.0, rel=0.0    )def test_local_loop_full_coordinate_length_random_cohort():    """Exercise the manifest coordinate length (L=1000) with a seeded cohort."""    length = 1000    n_contributors = 5    rng = random.Random(20260705)  # reproducible synthetic cohort    raw_vectors = [        [rng.choice((0, 1, 2)) for _ in range(length)] for _ in range(n_contributors)    ]    # Inject a couple of missing calls (null) to exercise the encode-as-0 path.    raw_vectors[0][3] = None    raw_vectors[2][17] = None    result = _run_pipeline(raw_vectors, length)    assert result["n_contributors"] == n_contributors  # sentinel == N, exactly    assert result["allele_counts"] == _cleartext_aggregate(raw_vectors, length)def test_sentinel_tracks_dropped_upload():    """Dropping one upload yields N-1 (the sentinel is a live contributor count)."""    length = 32    rng = random.Random(7)    raw_vectors = [        [rng.choice((0, 1, 2)) for _ in range(length)] for _ in range(6)    ]    full = _run_pipeline(raw_vectors, length)    dropped = _run_pipeline(raw_vectors[:-1], length)    assert full["n_contributors"] == 6    assert dropped["n_contributors"] == 5    # And the aggregate really lost exactly the dropped contributor's dosages.    last_encoded = local_data_owner.encode(raw_vectors[-1], length)    assert [        full_c - dropped_c        for full_c, dropped_c in zip(full["allele_counts"], dropped["allele_counts"])    ] == last_encoded@pytest.mark.parametrize("security", [128, 192, 256])def test_bit_exact_at_every_security_level(security):    """The full loop stays bit-exact vs the cleartext oracle at 128/192/256 bits,    and the sentinel decrypts to EXACTLY N — at each level. Also asserts the    coeff-modulus chain's ACHIEVED security equals the REQUESTED level."""    length = 128    n_contributors = 7    rng = random.Random(20260706 + security)    raw_vectors = [        [rng.choice((0, 1, 2)) for _ in range(length)] for _ in range(n_contributors)    ]    # A couple of missing calls to exercise the null->0 encode path under each chain.    raw_vectors[1][5] = None    raw_vectors[4][42] = None    result = _run_pipeline(raw_vectors, length, security=security)    # Sentinel recovers the EXACT contributor count at this security level.    assert result["n_contributors"] == n_contributors    # Encrypted aggregate == cleartext oracle, BIT-EXACT (BFV, tolerance 0).    assert result["allele_counts"] == _cleartext_aggregate(raw_vectors, length)    # The chain we shipped for this level actually certifies THIS level.    assert _achieved_security(8192, local_project_owner.SECURITY[security]) == security