tests/test_estimate.py

from src import config
from src.estimate import Estimate, format_estimate, parse_estimate


def test_parse_picks_last():
    logs = [
        "noise",
        "ESTIMATE n_records=10 total_bytes=100",
        "more",
        "ESTIMATE n_records=2 total_bytes=685",
    ]
    est = parse_estimate(logs, "cpu-upgrade")
    assert (est.n_records, est.total_bytes) == (2, 685)


def test_parse_none_when_absent():
    assert parse_estimate(["nothing here"], "cpu-upgrade") is None


def test_cost_math():
    # 1.5M records on cpu-upgrade @1500 rec/s -> 1000s fetch; multi-process -> merge.
    total_bytes = 3 * 1024 ** 3
    est = Estimate(n_records=1_500_000, total_bytes=total_bytes, flavor="cpu-upgrade")
    assert est.rec_per_s == config.FLAVOR_REC_PER_S["cpu-upgrade"]
    assert abs(est.fetch_seconds - 1000.0) < 1e-6
    assert est.n_processes == 8  # capped at cpu-upgrade vCPUs -> multi-process
    expected_merge = config.MERGE_BASE_S + total_bytes / config.MERGE_THROUGHPUT_BYTES_PER_S
    assert abs(est.merge_seconds - expected_merge) < 1e-6
    expected_duration = config.CONTAINER_OVERHEAD_S + 1000.0 + expected_merge
    assert abs(est.duration_seconds - expected_duration) < 1e-6
    expected_cost = expected_duration / 3600 * config.FLAVOR_HOURLY_USD["cpu-upgrade"]
    assert abs(est.cost_usd - expected_cost) < 1e-9
    assert est.n_output_files == 4  # 3 GiB ≈ 3.22 GB / 1 GB target -> ceil = 4
    md = format_estimate(est)
    assert "Cost estimate" in md and "merge" in md


def test_small_job_no_merge():
    # Few records -> single process -> no merge step.
    est = Estimate(n_records=2, total_bytes=685, flavor="cpu-upgrade")
    assert est.n_processes == 1
    assert est.merge_seconds == 0.0
    assert abs(est.duration_seconds - (config.CONTAINER_OVERHEAD_S + est.fetch_seconds)) < 1e-6
    assert "merge" not in format_estimate(est)


def test_zero_records_message():
    est = Estimate(0, 0, "cpu-upgrade")
    assert est.n_output_files == 0
    assert "No matching records" in format_estimate(est)


def test_cost_report():
    from src.estimate import actual_cost_usd, format_cost_report

    # 1 hour on cpu-upgrade ($0.03/hr) -> $0.03
    assert abs(actual_cost_usd(3600, "cpu-upgrade") - 0.03) < 1e-9
    r = format_cost_report(46, "cpu-upgrade")
    assert "Actual cost" in r and "46s" in r and "cpu-upgrade" in r


def test_unknown_flavor_falls_back():
    est = Estimate(1000, 1000, "mystery")
    assert est.rec_per_s == config.FLAVOR_REC_PER_S[config.DEFAULT_FLAVOR]
    assert est.hourly_usd == config.FLAVOR_HOURLY_USD[config.DEFAULT_FLAVOR]