gpu/single-gpu/home_benjamin_hdbasin/code/tests/test_data_mixture.py

import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
SCRIPTS = ROOT / "scripts"
if str(SCRIPTS) not in sys.path:
    sys.path.insert(0, str(SCRIPTS))

from run_data_mixture_benchmark import parse_args
from run_llm_lora_sft_benchmark import resolve_method_specs


def test_data_mixture_defaults_to_hdbasinflow_only():
    args = parse_args([])

    assert args.methods == "hdbasinflow"
    assert resolve_method_specs(args.methods) == [("HD-BasinFlow", "full", 20)]


def test_data_mixture_comparison_is_explicit_opt_in():
    args = parse_args(["--methods", "comparison"])
    specs = resolve_method_specs(args.methods)

    assert ("HD-BasinFlow", "full", 20) in specs
    assert ("IsoBasinFlow", "full", 100) in specs
    assert ("Random Search", "baseline", 1) in specs