test_json.py

import textwrap

import pyarrow as pa
import pytest

from datasets.packaged_modules.json.json import Json


@pytest.fixture
def jsonl_file(tmp_path):
    filename = tmp_path / "file.jsonl"
    data = textwrap.dedent(
        """\
        {"col_1": 1, "col_2": 2}
        {"col_1": 10, "col_2": 20}
        """
    )
    with open(filename, "w") as f:
        f.write(data)
    return str(filename)


@pytest.fixture
def json_file_with_list_of_dicts(tmp_path):
    filename = tmp_path / "file_with_list_of_dicts.json"
    data = textwrap.dedent(
        """\
        [
            {"col_1": 1, "col_2": 2},
            {"col_1": 10, "col_2": 20}
        ]
        """
    )
    with open(filename, "w") as f:
        f.write(data)
    return str(filename)


@pytest.fixture
def json_file_with_list_of_dicts_field(tmp_path):
    filename = tmp_path / "file_with_list_of_dicts_field.json"
    data = textwrap.dedent(
        """\
        {
            "field1": 1,
            "field2": "aabb",
            "field3": [
                {"col_1": 1, "col_2": 2},
                {"col_1": 10, "col_2": 20}
            ]
        }
        """
    )
    with open(filename, "w") as f:
        f.write(data)
    return str(filename)


@pytest.mark.parametrize(
    "file_fixture, config_kwargs",
    [
        ("jsonl_file", {}),
        ("json_file_with_list_of_dicts", {}),
        ("json_file_with_list_of_dicts_field", {"field": "field3"}),
    ],
)
def test_json_generate_tables(file_fixture, config_kwargs, request):
    json = Json(**config_kwargs)
    generator = json._generate_tables([[request.getfixturevalue(file_fixture)]])
    pa_table = pa.concat_tables([table for _, table in generator])
    assert pa_table.to_pydict() == {"col_1": [1, 10], "col_2": [2, 20]}