data_studio/audits.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from data_studio.sampling import SamplingManifestV1, reproducible_random_sample


@dataclass(frozen=True)
class AuditPolicy:
    policy_id: str
    random_sample_size: int
    maximum_random_failures: int
    risk_sample_size: int
    maximum_critical_failures: int = 0


AUDIT_POLICY_100K_DRAFT = AuditPolicy(
    policy_id="bulk-100k-draft-v1",
    random_sample_size=700,
    maximum_random_failures=14,
    risk_sample_size=300,
)


def build_audit_samples(
    population_ids: list[str],
    risk_ids: list[str],
    policy: AuditPolicy,
    seed: str,
) -> tuple[SamplingManifestV1, SamplingManifestV1]:
    random_manifest = reproducible_random_sample(population_ids, policy.random_sample_size, f"{seed}:random")
    risk_population = sorted(set(risk_ids).difference(random_manifest.selected_record_ids))
    risk_manifest = reproducible_random_sample(risk_population, policy.risk_sample_size, f"{seed}:risk")
    return random_manifest, risk_manifest


def audit_promotion_result(
    policy: AuditPolicy,
    random_result: dict[str, Any],
    risk_result: dict[str, Any],
) -> dict[str, Any]:
    random_failures = sum(
        int(random_result.get(key, 0))
        for key in ("minor_failure_count", "major_failure_count", "critical_failure_count")
    )
    critical = int(random_result.get("critical_failure_count", 0)) + int(
        risk_result.get("critical_failure_count", 0)
    )
    systematic_risk_failure = bool(risk_result.get("systematic_failure", False))
    passed = (
        random_failures <= policy.maximum_random_failures
        and critical <= policy.maximum_critical_failures
        and not systematic_risk_failure
    )
    return {
        "policy_id": policy.policy_id,
        "passed": passed,
        "random_failures": random_failures,
        "critical_failures": critical,
        "systematic_risk_failure": systematic_risk_failure,
    }