Repository: huggingface/datasets
File: tests/test_arrow_dataset.py (Python)

Code before:
    dataset = Dataset.from_generator(data_generator, features=features, cache_dir=cache_dir)
    _check_generator_dataset(dataset, expected_features, NamedSplit("train"))


@pytest.mark.parametrize(
    "split",
    [None, NamedSplit("train"), "train", NamedSplit("foo"), "foo"],
)
def test_dataset_from_generator_split(split, data_generator, tmp_path):
    cache_dir = tmp_path / "cache"
    default_expected_split = "train"
    expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
    expected_split = split if split else default_expected_split
    if split:
        dataset = Dataset.from_generator(data_generator, cache_dir=cache_dir, split=split)
    else:
        dataset = Dataset.from_generator(data_generator, cache_dir=cache_dir)
    _check_generator_dataset(dataset, expected_features, expected_split)


@pytest.mark.parametrize("fingerprint", [None, "test-dataset"])
def test_dataset_from_generator_fingerprint(fingerprint, data_generator, tmp_path):
    cache_dir = tmp_path / "cache"
    expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
    dataset = Dataset.from_generator(data_generator, cache_dir=cache_dir, fingerprint=fingerprint)
    _check_generator_dataset(dataset, expected_features, NamedSplit("train"))


@require_not_windows
@require_dill_gt_0_3_2
@require_pyspark
def test_from_spark():
    import pyspark

    spark = pyspark.sql.SparkSession.builder.master("local[*]").appName("pyspark").getOrCreate()
    data = [
        ("0", 0, 0.0),
        ("1", 1, 1.0),
        ("2", 2, 2.0),
        ("3", 3, 3.0),
    ]
    df = spark.createDataFrame(data, "col_1: string, col_2: int, col_3: float")
    dataset = Dataset.from_spark(df)
    assert isinstance(dataset, Dataset)
    assert dataset.num_rows == 4
    assert dataset.num_columns == 3
    assert dataset.column_names == ["col_1", "col_2", "col_3"]


@require_not_windows
@require_dill_gt_0_3_2

Review comment (suggestion):
added this test

```suggestion
    _check_generator_dataset(dataset, expected_features, NamedSplit("train"))
    if fingerprint:
        assert dataset._fingerprint == fingerprint
```

Code after:
    dataset = Dataset.from_generator(data_generator, features=features, cache_dir=cache_dir)
    _check_generator_dataset(dataset, expected_features, NamedSplit("train"))


@pytest.mark.parametrize(
    "split",
    [None, NamedSplit("train"), "train", NamedSplit("foo"), "foo"],
)
def test_dataset_from_generator_split(split, data_generator, tmp_path):
    cache_dir = tmp_path / "cache"
    default_expected_split = "train"
    expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
    expected_split = split if split else default_expected_split
    if split:
        dataset = Dataset.from_generator(data_generator, cache_dir=cache_dir, split=split)
    else:
        dataset = Dataset.from_generator(data_generator, cache_dir=cache_dir)
    _check_generator_dataset(dataset, expected_features, expected_split)


@pytest.mark.parametrize("fingerprint", [None, "test-dataset"])
def test_dataset_from_generator_fingerprint(fingerprint, data_generator, tmp_path):
    cache_dir = tmp_path / "cache"
    expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
    dataset = Dataset.from_generator(data_generator, cache_dir=cache_dir, fingerprint=fingerprint)
    _check_generator_dataset(dataset, expected_features, NamedSplit("train"))
    if fingerprint:
        assert dataset._fingerprint == fingerprint


@require_not_windows
@require_dill_gt_0_3_2
@require_pyspark
def test_from_spark():
    import pyspark

    spark = pyspark.sql.SparkSession.builder.master("local[*]").appName("pyspark").getOrCreate()
    data = [
        ("0", 0, 0.0),
        ("1", 1, 1.0),
        ("2", 2, 2.0),
        ("3", 3, 3.0),
    ]
    df = spark.createDataFrame(data, "col_1: string, col_2: int, col_3: float")
    dataset = Dataset.from_spark(df)
    assert isinstance(dataset, Dataset)
    assert dataset.num_rows == 4
    assert dataset.num_columns == 3
    assert dataset.column_names == ["col_1", "col_2", "col_3"]