#
# 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 that run Pipelines against a Spark Connect server.
"""

import unittest

from pyspark import pipelines as dp
from pyspark.testing.connectutils import (
    ReusedConnectTestCase,
    should_test_connect,
    connect_requirement_message,
)

if should_test_connect:
    from pyspark.errors.exceptions.connect import AnalysisException
    from pyspark.pipelines.graph_element_registry import graph_element_registration_context
    from pyspark.pipelines.spark_connect_graph_element_registry import (
        SparkConnectGraphElementRegistry,
    )
    from pyspark.pipelines.spark_connect_pipeline import (
        create_dataflow_graph,
        start_run,
        handle_pipeline_events,
    )


@unittest.skipIf(not should_test_connect, connect_requirement_message)
class SparkConnectPipelinesTest(ReusedConnectTestCase):
    def test_dry_run(self):
        dataflow_graph_id = create_dataflow_graph(self.spark, None, None, None)
        registry = SparkConnectGraphElementRegistry(self.spark, dataflow_graph_id)

        with graph_element_registration_context(registry):

            @dp.materialized_view
            def mv():
                return self.spark.range(1)

        result_iter = start_run(
            self.spark,
            dataflow_graph_id,
            full_refresh=None,
            refresh=None,
            full_refresh_all=False,
            dry=True,
            storage="file:///tmp/storage_path",
        )
        handle_pipeline_events(result_iter)

    def test_dry_run_failure(self):
        dataflow_graph_id = create_dataflow_graph(self.spark, None, None, None)
        registry = SparkConnectGraphElementRegistry(self.spark, dataflow_graph_id)

        with graph_element_registration_context(registry):

            @dp.table
            def st():
                # Invalid because a streaming query is expected
                return self.spark.range(1)

        result_iter = start_run(
            self.spark,
            dataflow_graph_id,
            full_refresh=None,
            refresh=None,
            full_refresh_all=False,
            dry=True,
            storage="file:///tmp/storage_path",
        )
        with self.assertRaises(AnalysisException) as context:
            handle_pipeline_events(result_iter)
        self.assertIn(
            "INVALID_FLOW_QUERY_TYPE.BATCH_RELATION_FOR_STREAMING_TABLE", str(context.exception)
        )


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

    main()
