#
# 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 logging
import unittest

import yaml

import apache_beam as beam
from apache_beam.yaml import YamlTransform
from apache_beam.yaml import yaml_provider
from apache_beam.yaml.yaml_provider import InlineProvider
from apache_beam.yaml.yaml_transform import Scope
from apache_beam.yaml.yaml_transform import chain_as_composite
from apache_beam.yaml.yaml_transform import ensure_errors_consumed
from apache_beam.yaml.yaml_transform import ensure_transforms_have_types
from apache_beam.yaml.yaml_transform import expand_composite_transform
from apache_beam.yaml.yaml_transform import expand_pipeline
from apache_beam.yaml.yaml_transform import extract_name
from apache_beam.yaml.yaml_transform import get_main_output_key
from apache_beam.yaml.yaml_transform import identify_object
from apache_beam.yaml.yaml_transform import normalize_inputs_outputs
from apache_beam.yaml.yaml_transform import normalize_source_sink
from apache_beam.yaml.yaml_transform import only_element
from apache_beam.yaml.yaml_transform import pipeline_as_composite
from apache_beam.yaml.yaml_transform import preprocess
from apache_beam.yaml.yaml_transform import preprocess_flattened_inputs
from apache_beam.yaml.yaml_transform import preprocess_windowing
from apache_beam.yaml.yaml_transform import push_windowing_to_roots
from apache_beam.yaml.yaml_transform import validate_against_schema
from apache_beam.yaml.yaml_utils import SafeLineLoader

try:
  import jsonschema
except ImportError:
  jsonschema = None


def new_pipeline():
  return beam.Pipeline(
      options=beam.options.pipeline_options.PipelineOptions(
          pickle_library='cloudpickle'))


@unittest.skipIf(jsonschema is None, "Yaml dependencies not installed")
class MainTest(unittest.TestCase):
  def assertYaml(self, expected, result):
    result = SafeLineLoader.strip_metadata(result)
    expected = yaml.load(expected, Loader=SafeLineLoader)
    expected = SafeLineLoader.strip_metadata(expected)

    self.assertEqual(expected, result)

  def get_scope_by_spec(self, p, spec, inputs=None):
    if inputs is None:
      inputs = {}
    spec = yaml.load(spec, Loader=SafeLineLoader)
    if 'transforms' in spec:
      spec['transforms'] = [
          normalize_inputs_outputs(t) for t in spec['transforms']
      ]

    scope = Scope(
        beam.pvalue.PBegin(p),
        inputs,
        spec['transforms'],
        yaml_provider.standard_providers(), {})
    return scope, spec

  def test_pipeline_as_composite_with_type_transforms(self):
    spec = '''
      type: composite
      transforms:
      - type: Create
        config:
          elements: [0,1,2]
      - type: PyMap
        config:
          fn: 'lambda x: x*x'
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    result = pipeline_as_composite(spec)

    self.assertEqual(result['type'], 'composite')
    self.assertEqual(result['name'], None)

  def test_pipeline_as_composite_with_transforms(self):
    spec = '''
      transforms:
      - type: Create
        config:
          elements: [0,1,2]
      - type: PyMap
        config:
          fn: 'lambda x: x*x'
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    result = pipeline_as_composite(spec)

    self.assertEqual(result['type'], 'composite')
    self.assertEqual(result['name'], None)

  def test_pipeline_as_composite_list(self):
    spec = '''
        - type: Create
          config:
            elements: [0,1,2]
        - type: PyMap
          config:
            fn: 'lambda x: x*x'
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    result = pipeline_as_composite(spec)

    expected = '''
      type: composite
      name: null
      transforms:
      - type: Create
        config:
          elements: [0,1,2]
      - type: PyMap
        config:
          fn: 'lambda x: x*x'
    '''
    self.assertYaml(expected, result)

  def test_expand_composite_transform_with_name(self):
    with new_pipeline() as p:
      spec = '''
        type: composite
        name: Custom
        transforms:
          - type: Create
            config:
              elements: [0,1,2]
        output:
          Create
        '''
      scope, spec = self.get_scope_by_spec(p, spec)
      self.assertRegex(
          str(expand_composite_transform(spec, scope)['output']),
          r"PCollection.*Custom/Create/Map.*")

  def test_expand_composite_transform_with_name_input(self):
    with new_pipeline() as p:
      spec = '''
        type: composite
        input: elements
        transforms:
          - type: LogForTesting
            input: input
        output:
          LogForTesting
        '''
      elements = p | beam.Create(range(3))
      scope, spec = self.get_scope_by_spec(p, spec,
                                           inputs={'elements': elements})
      self.assertRegex(
          str(expand_composite_transform(spec, scope)['output']),
          r"PCollection.*Composite/log_for_testing/LogForTesting.*")

  def test_expand_composite_transform_root(self):
    with new_pipeline() as p:
      spec = '''
        type: composite
        transforms:
          - type: Create
            config:
              elements: [0,1,2]
        output:
          Create
        '''
      scope, spec = self.get_scope_by_spec(p, spec)
      self.assertRegex(
          str(expand_composite_transform(spec, scope)['output']),
          r"PCollection.*Create/Map.*")

  def test_chain_as_composite(self):
    spec = '''
        type: chain
        transforms:
        - type: Create
          config:
            elements: [0,1,2]
        - type: PyMap
          config:
            fn: 'lambda x: x*x'
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    result = chain_as_composite(spec)

    expected = f'''
      type: composite
      name: Chain
      transforms:
      - type: Create
        config:
          elements: [0,1,2]
      - type: PyMap
        config:
          fn: 'lambda x: x*x'
        input: {spec['transforms'][0]['__uuid__']}
      output:
        '__implicit_outputs__': {spec['transforms'][1]['__uuid__']}
    '''
    self.assertYaml(expected, result)

  def test_chain_as_composite_with_wrong_output_type(self):
    spec = '''
        type: chain
        transforms:
        - type: Create
          elements: [0,1,2]
        - type: PyMap
          fn: 'lambda x: x*x'
        output:
          Create
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    with self.assertRaisesRegex(ValueError,
                                r"Explicit output.*of the chain transform is "
                                r"not an output of the last transform"):
      chain_as_composite(spec)

  def test_chain_as_composite_with_wrong_output_name(self):
    spec = '''
        type: chain
        transforms:
        - type: Create
          name: elements
          elements: [0,1,2]
        - type: PyMap
          fn: 'lambda x: x*x'
        output:
          elements
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    with self.assertRaisesRegex(ValueError,
                                r"Explicit output.*of the chain transform is "
                                r"not an output of the last transform"):
      chain_as_composite(spec)

  def test_chain_as_composite_with_outputs_override(self):
    spec = '''
        type: chain
        transforms:
        - type: Create
          elements: [0,1,2]
        - type: PyMap
          fn: 'lambda x: x*x'
        output:
          PyMap
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    self.assertEqual(
        chain_as_composite(spec)['output']['output'],
        f"{spec['transforms'][1]['__uuid__']}.PyMap")

  def test_chain_as_composite_with_input(self):
    spec = '''
        type: chain
        input:
          elements
        transforms:
        - type: PyMap
          fn: 'lambda x: x*x'
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    self.assertEqual(
        chain_as_composite(spec)['transforms'][0]['input'], {"input": "input"})

  def test_chain_as_composite_with_transform_input(self):
    spec = '''
        type: chain
        transforms:
        - type: Create
          config:
            elements: [0,1,2]
        - type: LogForTesting
          input: Create
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    with self.assertRaisesRegex(
        ValueError,
        r"Transform .* is part of a chain. "
        r"Cannot define explicit inputs on chain pipeline"):
      chain_as_composite(spec)

  def test_validate_against_schema_with_transform_input(self):
    spec = '''
        pipeline:
          type: chain
          transforms:
            - type: Create
              config:
                elements: [1, 2, 3]
            - type: LogForTesting
              input: Create
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    with self.assertRaisesRegex(jsonschema.ValidationError,
                                r"'input' should not be used "
                                r"along with 'chain' type transforms.*"):
      validate_against_schema(spec, 'generic')

  def test_normalize_source_sink(self):
    spec = '''
        source:
          type: Create
          config:
            elements: [0,1,2]
        transforms:
        - type: PyMap
          config:
            fn: 'lambda x: x*x'
        sink:
          type: PyMap
          config:
            fn: "lambda x: x + 41"
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    result = normalize_source_sink(spec)

    expected = '''
      transforms:
      - type: Create
        input: {'__explicitly_empty__': null}
        config:
          elements: [0,1,2]
      - type: PyMap
        config:
          fn: 'lambda x: x*x'
      - type: PyMap
        config:
          fn: "lambda x: x + 41"
    '''
    self.assertYaml(expected, result)

  def test_normalize_source_sink_only_source(self):
    spec = '''
        source:
          type: Create
          config:
            elements: [0,1,2]
        transforms:
        - type: PyMap
          config:
            fn: 'lambda x: x*x'

      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    result = normalize_source_sink(spec)

    expected = '''
      transforms:
      - type: Create
        input: {'__explicitly_empty__': null}
        config:
          elements: [0,1,2]
      - type: PyMap
        config:
          fn: 'lambda x: x*x'
    '''
    self.assertYaml(expected, result)

  def test_normalize_source_sink_only_sink(self):
    spec = '''
        transforms:
        - type: PyMap
          config:
            fn: 'lambda x: x*x'
        sink:
          type: PyMap
          config:
            fn: "lambda x: x + 41"
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    result = normalize_source_sink(spec)
    expected = '''
      transforms:
      - type: PyMap
        config:
          fn: 'lambda x: x*x'
      - type: PyMap
        config:
          fn: "lambda x: x + 41"
    '''
    self.assertYaml(expected, result)

  def test_normalize_source_sink_no_source_no_sink(self):
    spec = '''
        transforms:
        - type: PyMap
          config:
            fn: 'lambda x: x*x'
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    result = normalize_source_sink(spec)

    expected = '''
      transforms:
      - type: PyMap
        config:
          fn: 'lambda x: x*x'
    '''
    self.assertYaml(expected, result)

  def test_preprocess_source_sink_composite(self):
    spec = '''
      type: composite
      source:
        type: Create
        config:
          elements: [0,1,2]
      transforms:
      - type: PyMap
        config:
          fn: 'lambda x: x*x'
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    result = normalize_source_sink(spec)

    expected = '''
      type: composite
      transforms:
      - type: Create
        input: {'__explicitly_empty__': null}
        config:
          elements: [0,1,2]
      - type: PyMap
        config:
          fn: 'lambda x: x*x'
    '''
    self.assertYaml(expected, result)

  def test_preprocess_source_sink_chain(self):
    spec = '''
      type: chain
      source:
        type: Create
        config:
          elements: [0,1,2]
      transforms:
      - type: PyMap
        config:
          fn: 'lambda x: x*x'
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    result = normalize_source_sink(spec)

    expected = '''
      type: chain
      transforms:
      - type: Create
        input: {'__explicitly_empty__': null}
        config:
          elements: [0,1,2]
      - type: PyMap
        config:
          fn: 'lambda x: x*x'
    '''
    self.assertYaml(expected, result)

  def test_preprocess_source_sink_other(self):
    spec = '''
      - type: PyMap
        fn: 'lambda x: x*x'
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    self.assertEqual(normalize_source_sink(spec), spec)

  def test_normalize_inputs_outputs(self):
    spec = '''
        type: PyMap
        input: [Create1, Create2]
        fn: 'lambda x: x*x'
        output: Squared
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    result = normalize_inputs_outputs(spec)

    expected = '''
      type: PyMap
      input:
        input: [Create1, Create2]
      fn: 'lambda x: x*x'
      output:
        output: Squared
    '''
    self.assertYaml(expected, result)

  def test_normalize_inputs_outputs_dict(self):
    spec = '''
        type: PyMap
        input: [Create1, Create2]
        fn: 'lambda x: x*x'
        output:
          out1: Squared1
          out2: Squared2
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    result = normalize_inputs_outputs(spec)

    expected = '''
      type: PyMap
      input:
        input: [Create1, Create2]
      fn: 'lambda x: x*x'
      output:
        out1: Squared1
        out2: Squared2
    '''
    self.assertYaml(expected, result)

  def test_identify_object_with_name(self):
    spec = '''
      type: PyMap
      fn: 'lambda x: x*x'
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    self.assertRegex(identify_object(spec), r"PyMap.*[0-9]")

  def test_identify_object(self):
    spec = '''
      argument: PyMap
      fn: 'lambda x: x*x'
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    self.assertRegex(identify_object(spec), r"at.*[0-9]")

  def test_extract_name_by_type(self):
    spec = '''
      type: PyMap
      fn: 'lambda x: x*x'
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    self.assertEqual(extract_name(spec), "PyMap")

  def test_extract_name_by_id(self):
    spec = '''
      type: PyMap
      id: PyMapId
      fn: 'lambda x: x*x'
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    self.assertEqual(extract_name(spec), "PyMapId")

  def test_extract_name_by_name(self):
    spec = '''
      type: PyMap
      name: PyMapName
      fn: 'lambda x: x*x'
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    self.assertEqual(extract_name(spec), "PyMapName")

  def test_extract_name_no_name(self):
    spec = '''
      transforms:
      - arg: PyMap
        fn: 'lambda x: x*x'
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    self.assertEqual(extract_name(spec), "")

  def test_push_windowing_to_roots(self):
    spec = '''
      type: composite
      transforms:
      - type: Create
        elements: [0,1,2]
      - type: PyMap
        fn: 'lambda x: x*x'
        input: Create
      windowing:
        type: fixed
        size: 2
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    spec = normalize_inputs_outputs(spec)
    spec['transforms'] = [
        normalize_inputs_outputs(t) for t in spec['transforms']
    ]

    result = push_windowing_to_roots(spec)

    expected = '''
      type: composite
      transforms:
      - type: Create
        elements: [0,1,2]
        windowing:
          type: fixed
          size: 2
        __consumed_outputs:
          - null
        input: {}
        output: {}
      - type: PyMap
        fn: 'lambda x: x*x'
        input:
          input: Create
        output: {}
      windowing:
        type: fixed
        size: 2
      input: {}
      output: {}
    '''
    self.assertYaml(expected, result)

  def test_preprocess_windowing_custom_type(self):
    spec = '''
        type: SumGlobally
        input: Create
        windowing:
          type: fixed
          size: 4
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    spec = normalize_inputs_outputs(spec)
    result = preprocess_windowing(spec)

    expected = f'''
      type: composite
      name: SumGlobally
      input:
        input: Create
      transforms:
        - type: SumGlobally
          input:
            input: {result['transforms'][1]['__uuid__']}
          output: {{}}
        - type: WindowInto
          name: WindowInto[input]
          windowing:
            type: fixed
            size: 4
          input: {{input: input}}
      output: {result['transforms'][0]['__uuid__']}
      config: 
        error_handling: {{}}
    '''
    self.assertYaml(expected, result)

  def test_preprocess_windowing_composite_with_windowing_outer(self):
    spec = '''
      type: composite
      transforms:
        - type: CreateTimestamped
          name: Create
          elements: [0, 2, 4]
        - type: SumGlobally
          input: Create
      windowing:
        type: fixed
        size: 4
      output: SumGlobally
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    spec = normalize_inputs_outputs(spec)
    spec['transforms'] = [
        normalize_inputs_outputs(t) for t in spec['transforms']
    ]

    result = preprocess_windowing(spec)

    expected = '''
      type: composite
      input: {}
      transforms:
        - type: CreateTimestamped
          name: Create
          elements: [0, 2, 4]
          windowing:
            type: fixed
            size: 4
          __consumed_outputs:
            - null
          input: {}
          output: {}
        - type: SumGlobally
          input:
            input: Create
          output: {}
      output:
        output: SumGlobally
    '''
    self.assertYaml(expected, result)

  def test_preprocess_windowing_composite_with_windowing_on_input(self):
    spec = '''
      type: composite
      transforms:
        - type: CreateTimestamped
          name: Create
          elements: [0, 2, 4]
        - type: SumGlobally
          input: Create
          windowing:
            type: fixed
            size: 4
      output: SumGlobally
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    spec = normalize_inputs_outputs(spec)
    spec['transforms'] = [
        normalize_inputs_outputs(t) for t in spec['transforms']
    ]

    result = preprocess_windowing(spec)

    expected = '''
      type: composite
      input: {}
      transforms:
        - type: CreateTimestamped
          name: Create
          elements: [0, 2, 4]
          input: {}
          output: {}
        - type: SumGlobally
          input:
            input: Create
          windowing:
            type: fixed
            size: 4
          output: {}
      output:
        output: SumGlobally
    '''
    self.assertYaml(expected, result)

  def test_preprocess_windowing_other_type_with_no_inputs(self):
    spec = '''
      type: CreateTimestamped
      name: Create
      elements: [0, 2, 4]
      windowing:
        type: fixed
        size: 4
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    spec = normalize_inputs_outputs(spec)
    result = preprocess_windowing(spec)

    expected = f'''
      type: composite
      name: Create
      transforms:
        - type: CreateTimestamped
          name: Create
          elements: [0, 2, 4]
          input: {{}}
          output: {{}}
        - type: WindowInto
          name: WindowInto[None]
          input:
            input: {result['transforms'][0]["__uuid__"]}
          windowing:
            type: fixed
            size: 4
      output: {result['transforms'][1]["__uuid__"]}
      config: 
        error_handling: {{}}
    '''
    self.maxDiff = 1e9

    self.assertYaml(expected, result)

  def test_preprocess_flattened_inputs_implicit(self):
    spec = '''
      type: composite
      transforms:
        - type: PyMap
          fn: 'lambda x: x*x'
          input: [Create1, Create2]
      output: CreateTimestamped
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    spec['transforms'] = [
        normalize_inputs_outputs(t) for t in spec['transforms']
    ]
    result = preprocess_flattened_inputs(spec)

    expected = f'''
      type: composite
      transforms:
        - type: Flatten
          name: PyMap-Flatten[input]
          input:
            input0: Create1
            input1: Create2
        - type: PyMap
          fn: 'lambda x: x*x'
          input:
            input: {result['transforms'][0]['__uuid__']}
          output: {{}}
      output: CreateTimestamped
    '''
    self.assertYaml(expected, result)

  def test_preprocess_flattened_inputs_explicit_flatten(self):
    spec = '''
      type: composite
      transforms:
        - type: Flatten
          input: [Create1, Create2]
        - type: PyMap
          fn: 'lambda x: x*x'
          input: Flatten
      output: CreateTimestamped
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    spec['transforms'] = [
        normalize_inputs_outputs(t) for t in spec['transforms']
    ]
    result = preprocess_flattened_inputs(spec)

    expected = '''
      type: composite
      transforms:
        - type: Flatten
          input:
            input0: Create1
            input1: Create2
          output: {}
        - type: PyMap
          fn: 'lambda x: x*x'
          input:
            input: Flatten
          output: {}
      output: CreateTimestamped
    '''
    self.assertYaml(expected, result)

  def test_ensure_transforms_have_types(self):
    spec = '''
      type: PyMap
      fn: 'lambda x: x*x'
      input: Flatten
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    result = ensure_transforms_have_types(spec)
    self.assertEqual(result, spec)

  def test_ensure_transforms_have_types_error(self):
    spec = '''
      name: PyMap
      fn: 'lambda x: x*x'
      input: Flatten
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    with self.assertRaisesRegex(ValueError, r"Missing type .*"):
      ensure_transforms_have_types(spec)
    with self.assertRaisesRegex(ValueError, r"Missing type .*"):
      preprocess(spec)

  def test_ensure_transforms_have_providers_error(self):
    spec = '''
      type: UnknownType
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    with self.assertRaisesRegex(ValueError,
                                r"Unknown type or missing provider .*"):
      preprocess(spec, known_transforms=['KnownType'])

  def test_ensure_errors_consumed_unconsumed(self):
    spec = '''
      type: composite
      transforms:
        - type: Create
          elements: [1,2,3]
        - type: MyTransform
          input: Create
          error_handling:
            output: errors
      output:
        good: MyTransform
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    spec = normalize_inputs_outputs(spec)
    spec['transforms'] = [
        normalize_inputs_outputs(t) for t in spec['transforms']
    ]
    with self.assertRaisesRegex(ValueError,
                                r"Unconsumed error.*MyTransform.errors"):
      ensure_errors_consumed(spec)

  def test_ensure_errors_consumed_in_transform(self):
    spec = '''
      type: composite
      transforms:
        - type: Create
          elements: [1,2,3]
        - type: MyTransform
          input: Create
          error_handling:
            output: errors
        - name: SaveToFile
          type: PyMap
          input: MyTransform.errors
      output:
        good: MyTransform
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    spec = normalize_inputs_outputs(spec)
    spec['transforms'] = [
        normalize_inputs_outputs(t) for t in spec['transforms']
    ]
    result = ensure_errors_consumed(spec)
    self.assertEqual(result, spec)
    self.assertEqual(result['transforms'], spec['transforms'])

  def test_ensure_errors_consumed_no_output_in_error_handling(self):
    spec = '''
      type: composite
      transforms:
        - type: Create
          elements: [1,2,3]
        - type: MyTransform
          input: Create
          error_handling:
            arg: errors
        - name: SaveToFile
          type: PyMap
          input: MyTransform.errors
      output:
        good: MyTransform
    '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    spec = normalize_inputs_outputs(spec)
    spec['transforms'] = [
        normalize_inputs_outputs(t) for t in spec['transforms']
    ]
    with self.assertRaisesRegex(ValueError, r"Missing output.*"):
      ensure_errors_consumed(spec)

  def test_only_element(self):
    self.assertEqual(only_element((1, )), 1)

  def test_get_main_output_key(self):
    spec = {'type': 'TestTransform'}
    error_handling_spec = {'output': 'invalid_rows'}

    # Case 1: 'output' key exists
    outputs = {'output': 1, 'another': 2}
    with self.assertLogs('apache_beam.yaml.yaml_transform',
                         level='WARNING') as al:
      self.assertEqual(
          get_main_output_key(spec, outputs, error_handling_spec), 'output')
      self.assertIn("Only the main output will be validated.", al.output[0])

    # Case 2: 'good' key exists, 'output' does not
    outputs = {'good': 1, 'another': 2}
    with self.assertLogs('apache_beam.yaml.yaml_transform',
                         level='WARNING') as al:
      self.assertEqual(
          get_main_output_key(spec, outputs, error_handling_spec), 'good')
      self.assertIn("Only the main output will be validated.", al.output[0])

    # Case 3: Only one output
    outputs = {'single_output': 1}
    self.assertEqual(
        get_main_output_key(spec, outputs, error_handling_spec),
        'single_output')

    # Case 4: Multiple outputs, no 'output' or 'good'
    outputs = {'another': 1, 'yet_another': 2}
    with self.assertRaisesRegex(
        ValueError, "Transform .* has outputs .* but none are named 'output'"):
      get_main_output_key(spec, outputs, error_handling_spec)

    # Case 5: Empty outputs
    outputs = {}
    with self.assertRaisesRegex(
        ValueError, "Transform .* has outputs .* but none are named 'output'"):
      get_main_output_key(spec, outputs, error_handling_spec)

    # Case 6: More than two outputs with 'good' present
    outputs = {'good': 1, 'bad': 2, 'something': 3}
    with self.assertLogs('apache_beam.yaml.yaml_transform',
                         level='WARNING') as al:
      self.assertEqual(
          get_main_output_key(spec, outputs, error_handling_spec), 'good')
      self.assertIn("Only the main output will be validated.", al.output[0])


class YamlTransformTest(unittest.TestCase):
  def test_init_with_string(self):
    provider1 = InlineProvider({"MyTransform1": lambda: beam.Map(lambda x: x)})
    provider2 = InlineProvider({"MyTransform2": lambda: beam.Map(lambda x: x)})

    providers_dict = {"p1": [provider1], "p2": [provider2]}

    spec = '''
        type: chain
        transforms:
          - type: Create
            elements: [1,2,3]
          - type: LogForTesting
      '''
    result = YamlTransform(spec, providers_dict)
    self.assertIn('p1', result._providers)  # check for custom providers
    self.assertIn('p2', result._providers)  # check for custom providers
    self.assertIn(
        'LogForTesting', result._providers)  # check for standard provider
    self.assertEqual(result._spec['type'], "composite")  # preprocessed spec

  def test_init_with_dict(self):
    spec = '''
        type: chain
        transforms:
          - type: Create
            config:
              elements: [1,2,3]
          - type: LogForTesting
      '''
    spec = yaml.load(spec, Loader=SafeLineLoader)
    result = YamlTransform(spec)
    self.assertIn(
        'LogForTesting', result._providers)  # check for standard provider
    self.assertEqual(result._spec['type'], "composite")  # preprocessed spec


class ExpandPipelineTest(unittest.TestCase):
  def test_expand_pipeline_with_pipeline_key_only(self):
    spec = '''
      pipeline:
        type: chain
        transforms:
          - type: Create
            config:
              elements: [1,2,3]
          - type: LogForTesting
    '''
    with new_pipeline() as p:
      expand_pipeline(p, spec, validate_schema=None)

  def test_expand_pipeline_with_pipeline_and_option_keys(self):
    spec = '''
      pipeline:
        type: chain
        transforms:
          - type: Create
            config:
              elements: [1,2,3]
          - type: LogForTesting
      options:
        streaming: false
    '''
    with new_pipeline() as p:
      expand_pipeline(p, spec, validate_schema=None)

  def test_expand_pipeline_with_extra_top_level_keys(self):
    spec = '''
      template:
        version: "1.0"
        author: "test_user"

      pipeline:
        type: chain
        transforms:
          - type: Create
            config:
              elements: [1,2,3]
          - type: LogForTesting

      other_metadata: "This is an ignored comment."
    '''
    with new_pipeline() as p:
      expand_pipeline(p, spec, validate_schema=None)

  def test_expand_pipeline_with_incorrect_pipelines_key_fails(self):
    spec = '''
      pipelines:
        type: chain
        transforms:
          - type: Create
            config:
              elements: [1,2,3]
          - type: LogForTesting
    '''
    with new_pipeline() as p:
      with self.assertRaises(KeyError):
        expand_pipeline(p, spec, validate_schema=None)

  @unittest.skipIf(jsonschema is None, "Yaml dependencies not installed")
  def test_expand_pipeline_with_valid_schema(self):
    spec = '''
      pipeline:
        type: chain
        transforms:
          - type: Create
            config:
              elements: [1,2,3]
          - type: LogForTesting
    '''
    with new_pipeline() as p:
      expand_pipeline(p, spec, validate_schema='generic')

  @unittest.skipIf(jsonschema is None, "Yaml dependencies not installed")
  def test_expand_pipeline_with_invalid_schema(self):
    spec = '''
      pipeline:
        type: chain
        transforms:
          - name: Create
            config:
              elements: [1,2,3]
          - type: LogForTesting
    '''
    with new_pipeline() as p:
      with self.assertRaises(jsonschema.ValidationError):
        expand_pipeline(p, spec, validate_schema='generic')


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  unittest.main()
