from pathlib import Path

from schemas_hf.cache import DatasetConfigRef, iter_dataset_cache_configs, iter_dataset_info_files


def test_iter_dataset_info_files(tmp_path):
    repo_dir = tmp_path / "datasets--owner--name"
    info_path = repo_dir / "snapshots" / "123abc" / "dataset_info.json"
    info_path.parent.mkdir(parents=True)
    info_path.write_text("{}")

    second_dir = tmp_path / "datasets--ag_news"
    second_path = second_dir / "dataset_info.json"
    second_path.parent.mkdir(parents=True)
    second_path.write_text("{}")

    files = sorted(iter_dataset_info_files(tmp_path), key=lambda f: f.repo_id)
    assert len(files) == 2
    assert files[0].repo_id == "ag_news"
    assert files[0].path == second_path
    assert files[1].repo_id == "owner/name"
    assert files[1].path == info_path


def test_iter_dataset_cache_configs(tmp_path):
    cache_dir = tmp_path / "datasets"
    (cache_dir / "ag_news" / "default").mkdir(parents=True)
    (cache_dir / "owner___name" / "cfg").mkdir(parents=True)
    (cache_dir / "downloads").mkdir()

    refs = sorted(
        iter_dataset_cache_configs(cache_dir),
        key=lambda ref: (ref.dataset_id, ref.config_name or ""),
    )

    assert refs == [
        DatasetConfigRef(dataset_id="ag_news", config_name=None),
        DatasetConfigRef(dataset_id="owner/name", config_name="cfg"),
    ]
