python/pyspark/pipelines/tests/test_init_cli.py - master - githubmirror/spark

python/pyspark/pipelines/tests/test_init_cli.py
93 строки4 KB
yangjie01
[SPARK-55986][PYTHON] Upgrade black to 26.3.1
16 мар 2026, 19:20
16 мар 2026, 19:20cbcee8c
#
# 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.
#
import unittest
import tempfile
from pathlib import Path

from pyspark.testing.connectutils import (
    ReusedConnectTestCase,
    should_test_connect,
    connect_requirement_message,
)
from pyspark.testing.utils import have_yaml, yaml_requirement_message

if should_test_connect and have_yaml:
    from pyspark.pipelines.cli import (
        change_dir,
        find_pipeline_spec,
        load_pipeline_spec,
        init,
        register_definitions,
    )
    from pyspark.pipelines.tests.local_graph_element_registry import LocalGraphElementRegistry


@unittest.skipIf(
    not should_test_connect or not have_yaml,
    connect_requirement_message or yaml_requirement_message,
)
class InitCLITests(ReusedConnectTestCase):
    def test_init(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            project_name = "test_project"
            with change_dir(Path(temp_dir)):
                init(project_name)
            with change_dir(Path(temp_dir) / project_name):
                spec_path = find_pipeline_spec(Path.cwd())
                spec = load_pipeline_spec(spec_path)
                assert spec.name == project_name

                # Verify that the storage path is an absolute URI with file scheme
                expected_storage_path = f"file://{Path.cwd() / 'pipeline-storage'}"
                self.assertEqual(spec.storage, expected_storage_path)

                # Verify that the storage directory was created
                self.assertTrue((Path.cwd() / "pipeline-storage").exists())

                registry = LocalGraphElementRegistry()
                register_definitions(spec_path, registry, spec, self.spark, "test_graph_id")
                self.assertEqual(len(registry.outputs), 1)
                self.assertEqual(registry.outputs[0].name, "example_python_materialized_view")
                self.assertEqual(len(registry.flows), 1)
                self.assertEqual(registry.flows[0].name, "example_python_materialized_view")
                self.assertEqual(registry.flows[0].target, "example_python_materialized_view")
                self.assertEqual(len(registry.sql_files), 1)
                self.assertEqual(
                    registry.sql_files[0].file_path,
                    Path("transformations") / "example_sql_materialized_view.sql",
                )

    def test_init_existing_directory(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            project_name = "test_project"
            with change_dir(Path(temp_dir)):
                init(project_name)

                with self.assertRaises(FileExistsError) as context:
                    init(project_name)

                expected_message = (
                    f"Directory '{project_name}' already exists. "
                    "Please choose a different name or remove the existing directory."
                )
                self.assertEqual(str(context.exception), expected_message)


if __name__ == "__main__":
    from pyspark.testing import main

    main()