<gh_stars>1-10
# Copyright 2017 StreamSets Inc.
#
# Licensed 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 os
import string
import tempfile
from datetime import datetime
from decimal import Decimal

from streamsets.testframework.markers import sdc_min_version
from streamsets.testframework.utils import get_random_string

logger = logging.getLogger(__name__)


@sdc_min_version('2.7.1.0')
def test_schema_generator_processor(sdc_builder, sdc_executor):
    """Schema generated by the schema generator processor is well exercised by unit tests, so this integration
    test more focuses on interoperability with our Avro DataFormat library. We create a rather complex record
    with all different data types, generate a schema for it, write it to an Avro file using the generated schema,
    and read it back in a second pipeline. We assert that what we read back (using the generated schema) is
    what we had at the beginning.

    Schema Generator Pipeline:
        dev_raw_data_source >> record_deduplicator >> field_type_converter >> schema_generator >> local_fs
                                                   >> to_error

    Reader Pipeline:
        directory >> trash
    """
    # Test write directory
    tmp_directory = os.path.join(tempfile.gettempdir(), get_random_string(string.ascii_letters, 10))

    # Build pipeline that will generate test record and it's schema
    builder = sdc_builder.get_pipeline_builder()

    dev_raw_data_source = builder.add_stage('Dev Raw Data Source')
    dev_raw_data_source.data_format = 'JSON'
    dev_raw_data_source.raw_data = """{
      "str" : "str",
      "int": 10,
      "boolean": true,
      "decimal": "10.5",
      "date": "2017-01-01",
      "time": "10:09:08",
      "datetime": "2017-01-01 10:09:08",
      "list": ["a", "b"],
      "map": {"first-key" : "value", "second-key" : "secret value"}
    }
    """

    record_deduplicator = builder.add_stage('Record Deduplicator')
    to_error = builder.add_stage('To Error')

    # We can't express all types in JSON directly, so we'll convert them explicitly
    field_type_converter = builder.add_stage('Field Type Converter')
    field_type_converter.conversion_method = 'BY_FIELD'
    field_type_converter.field_type_converter_configs = [{'fields': ['/decimal'],
                                                          'targetType': 'DECIMAL',
                                                          'dataLocale': 'en,US',
                                                          'scale': 2,
                                                          'decimalScaleRoundingStrategy': 'ROUND_UNNECESSARY'},
                                                         {'fields': ['/date'],
                                                          'targetType': 'DATE',
                                                          'dateFormat': 'YYYY_MM_DD'},
                                                         {'fields': ['/time'],
                                                          'targetType': 'TIME',
                                                          'dateFormat': 'OTHER',
                                                          'otherDateFormat': 'HH:mm:ss'},
                                                         {'fields': ['/datetime'],
                                                          'targetType': 'DATETIME',
                                                          'dateFormat': 'YYYY_MM_DD_HH_MM_SS'}]

    # Generate schema for that record
    schema_generator = builder.add_stage('Schema Generator')
    schema_generator.schema_name = 'test_schema'

    # And store it in local file system
    local_fs = builder.add_stage('Local FS', type='destination')
    local_fs.directory_template = tmp_directory
    local_fs.data_format = 'AVRO'
    local_fs.configuration['configs.dataGeneratorFormatConfig.avroSchemaSource'] = 'HEADER'

    # Finish building the pipeline
    dev_raw_data_source >> record_deduplicator >> field_type_converter >> schema_generator >> local_fs
    record_deduplicator >> to_error
    generator_pipeline = builder.build('Schema Generator Pipeline')

    # Build second pipeline that will read generated Avro file
    builder = sdc_builder.get_pipeline_builder()

    directory = builder.add_stage('Directory', type='origin')
    directory.data_format = 'AVRO'
    directory.batch_wait_time_in_secs = 1
    directory.file_name_pattern = 'sdc*'
    directory.files_directory = tmp_directory

    trash = builder.add_stage('Trash')

    directory >> trash
    reader_pipeline = builder.build('Reader pipeline')

    sdc_executor.add_pipeline(generator_pipeline, reader_pipeline)
    sdc_executor.start_pipeline(generator_pipeline).wait_for_pipeline_batch_count(1)
    sdc_executor.stop_pipeline(generator_pipeline)

    # Read written data with second pipeline
    snapshot = sdc_executor.capture_snapshot(reader_pipeline, start_pipeline=True).snapshot
    sdc_executor.stop_pipeline(reader_pipeline)

    # Validate the input record's structure
    assert len(snapshot.snapshot_batches[0][directory.instance_name].output) == 1
    record = snapshot.snapshot_batches[0][directory.instance_name].output[0]

    # Assert proper values & types
    assert 'str' == record.field['str'].value
    assert 10 == record.field['int'].value
    assert record.field['boolean'].value == True
    assert Decimal('10.50') == record.field['decimal'].value
    assert datetime(2017, 1, 1, 0, 0) == record.field['date'].value
    assert datetime(1970, 1, 1, 10, 9, 8) == record.field['time'].value
    assert datetime(2017, 1, 1, 10, 9, 8) == record.field['datetime'].value
    assert 'a' == record.field['list'][0].value
    assert 'b' == record.field['list'][1].value
    assert 'value' == record.field['map']['first-key'].value
    assert 'secret value' == record.field['map']['second-key'].value
