superset2/tests/unit_tests/examples/data_loading_test.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Tests for data_loading.py UUID extraction functionality."""

from pathlib import Path
from tempfile import TemporaryDirectory
from unittest.mock import patch

import yaml


def test_get_dataset_config_from_yaml_extracts_uuid():
    """Test that UUID is extracted from dataset.yaml."""
    from superset.examples.data_loading import get_dataset_config_from_yaml

    with TemporaryDirectory() as tmpdir:
        example_dir = Path(tmpdir)
        dataset_yaml = example_dir / "dataset.yaml"
        dataset_yaml.write_text(
            yaml.dump(
                {
                    "table_name": "test_table",
                    "uuid": "12345678-1234-1234-1234-123456789012",
                    "schema": "public",
                }
            )
        )

        config = get_dataset_config_from_yaml(example_dir)

        assert config["uuid"] == "12345678-1234-1234-1234-123456789012"
        assert config["table_name"] == "test_table"
        assert config["schema"] == "public"


def test_get_dataset_config_from_yaml_without_uuid():
    """Test that missing UUID returns None."""
    from superset.examples.data_loading import get_dataset_config_from_yaml

    with TemporaryDirectory() as tmpdir:
        example_dir = Path(tmpdir)
        dataset_yaml = example_dir / "dataset.yaml"
        dataset_yaml.write_text(
            yaml.dump(
                {
                    "table_name": "test_table",
                    "schema": "public",
                }
            )
        )

        config = get_dataset_config_from_yaml(example_dir)

        assert config["uuid"] is None
        assert config["table_name"] == "test_table"


def test_get_dataset_config_from_yaml_no_file():
    """Test behavior when dataset.yaml doesn't exist."""
    from superset.examples.data_loading import get_dataset_config_from_yaml

    with TemporaryDirectory() as tmpdir:
        example_dir = Path(tmpdir)

        config = get_dataset_config_from_yaml(example_dir)

        assert config["uuid"] is None
        assert config["table_name"] is None
        assert config["schema"] is None


def test_get_dataset_config_from_yaml_treats_main_schema_as_none():
    """Test that SQLite's 'main' schema is treated as None."""
    from superset.examples.data_loading import get_dataset_config_from_yaml

    with TemporaryDirectory() as tmpdir:
        example_dir = Path(tmpdir)
        dataset_yaml = example_dir / "dataset.yaml"
        dataset_yaml.write_text(
            yaml.dump(
                {
                    "table_name": "test_table",
                    "schema": "main",  # SQLite default schema
                }
            )
        )

        config = get_dataset_config_from_yaml(example_dir)

        assert config["schema"] is None


def test_get_multi_dataset_config_extracts_uuid():
    """Test that UUID is extracted from datasets/{name}.yaml."""
    from superset.examples.data_loading import _get_multi_dataset_config

    with TemporaryDirectory() as tmpdir:
        example_dir = Path(tmpdir)
        datasets_dir = example_dir / "datasets"
        datasets_dir.mkdir()
        dataset_yaml = datasets_dir / "test_dataset.yaml"
        dataset_yaml.write_text(
            yaml.dump(
                {
                    "table_name": "custom_table_name",
                    "uuid": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee",
                    "schema": "public",
                }
            )
        )

        data_file = example_dir / "data" / "test_dataset.parquet"
        config = _get_multi_dataset_config(example_dir, "test_dataset", data_file)

        assert config["uuid"] == "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee"
        assert config["table_name"] == "custom_table_name"


def test_get_multi_dataset_config_without_yaml():
    """Test behavior when datasets/{name}.yaml doesn't exist."""
    from superset.examples.data_loading import _get_multi_dataset_config

    with TemporaryDirectory() as tmpdir:
        example_dir = Path(tmpdir)
        data_file = example_dir / "data" / "test_dataset.parquet"

        config = _get_multi_dataset_config(example_dir, "test_dataset", data_file)

        assert config.get("uuid") is None
        assert config["table_name"] == "test_dataset"


def test_get_multi_dataset_config_treats_main_schema_as_none():
    """Test that SQLite's 'main' schema is treated as None in multi-dataset config."""
    from superset.examples.data_loading import _get_multi_dataset_config

    with TemporaryDirectory() as tmpdir:
        example_dir = Path(tmpdir)
        datasets_dir = example_dir / "datasets"
        datasets_dir.mkdir()
        dataset_yaml = datasets_dir / "test_dataset.yaml"
        dataset_yaml.write_text(
            yaml.dump(
                {
                    "table_name": "test_table",
                    "schema": "main",
                }
            )
        )

        data_file = example_dir / "data" / "test_dataset.parquet"
        config = _get_multi_dataset_config(example_dir, "test_dataset", data_file)

        assert config["schema"] is None


def test_discover_datasets_passes_uuid_to_loader():
    """Test that discover_datasets passes UUID from YAML to create_generic_loader."""
    from superset.examples.data_loading import discover_datasets

    with TemporaryDirectory() as tmpdir:
        examples_dir = Path(tmpdir)

        # Create a simple example with data.parquet and dataset.yaml
        example_dir = examples_dir / "test_example"
        example_dir.mkdir()
        (example_dir / "data.parquet").touch()
        (example_dir / "dataset.yaml").write_text(
            yaml.dump(
                {
                    "table_name": "test_table",
                    "uuid": "12345678-1234-1234-1234-123456789012",
                }
            )
        )

        with patch(
            "superset.examples.data_loading.get_examples_directory",
            return_value=examples_dir,
        ):
            with patch(
                "superset.examples.data_loading.create_generic_loader"
            ) as mock_create:
                mock_create.return_value = lambda: None

                discover_datasets()

                mock_create.assert_called_once()
                call_kwargs = mock_create.call_args[1]
                assert call_kwargs["uuid"] == "12345678-1234-1234-1234-123456789012"