#

# 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 textio module."""

# pytype: skip-file

import bz2

import glob

import gzip

import logging

import os

import platform

import re

import shutil

import tempfile

import unittest

import zlib

from datetime import datetime

import pytz

import apache_beam as beam

from apache_beam import coders

from apache_beam.io import iobase

from apache_beam.io import source_test_utils

from apache_beam.io.filesystem import CompressionTypes

from apache_beam.io.filesystems import FileSystems

# Importing following private classes for testing.

from apache_beam.io.textio import ReadAllFromText

from apache_beam.io.textio import ReadAllFromTextContinuously

from apache_beam.io.textio import ReadFromText

from apache_beam.io.textio import ReadFromTextWithFilename

from apache_beam.io.textio import WriteToText

from apache_beam.io.textio import _TextSink as TextSink

from apache_beam.io.textio import _TextSource as TextSource

from apache_beam.options.pipeline_options import PipelineOptions

from apache_beam.testing.test_pipeline import TestPipeline

from apache_beam.testing.test_stream import TestStream

from apache_beam.testing.test_utils import TempDir

from apache_beam.testing.util import assert_that

from apache_beam.testing.util import equal_to

from apache_beam.transforms.core import Create

from apache_beam.transforms.userstate import CombiningValueStateSpec

from apache_beam.transforms.util import LogElements

from apache_beam.utils.timestamp import Timestamp

class DummyCoder(coders.Coder):

def encode(self, x):

raise ValueError

def decode(self, x):

return (x * 2).decode('utf-8')

def to_type_hint(self):

return str

class EOL(object):

LF = 1

CRLF = 2

MIXED = 3

LF_WITH_NOTHING_AT_LAST_LINE = 4

CUSTOM_DELIMITER = 5

def write_data(

num_lines,

no_data=False,

directory=None,

prefix=tempfile.template,

eol=EOL.LF,

custom_delimiter=None,

line_value=b'line'):

"""Writes test data to a temporary file.

Args:

num_lines (int): The number of lines to write.

no_data (bool): If :data:`True`, empty lines will be written, otherwise

each line will contain a concatenation of b'line' and the line number.

directory (str): The name of the directory to create the temporary file in.

prefix (str): The prefix to use for the temporary file.

eol (int): The line ending to use when writing.

:class:`~apache_beam.io.textio_test.EOL` exposes attributes that can be

used here to define the eol.

custom_delimiter (bytes): The custom delimiter.

line_value (bytes): Default value for test data, default b'line'

Returns:

Tuple[str, List[str]]: A tuple of the filename and a list of the

utf-8 decoded written data.

"""

all_data = []

with tempfile.NamedTemporaryFile(delete=False, dir=directory,

prefix=prefix) as f:

sep_values = [b'\n', b'\r\n']

for i in range(num_lines):

data = b'' if no_data else line_value + str(i).encode()

all_data.append(data)

if eol == EOL.LF:

sep = sep_values[0]

elif eol == EOL.CRLF:

sep = sep_values[1]

elif eol == EOL.MIXED:

sep = sep_values[i % len(sep_values)]

elif eol == EOL.LF_WITH_NOTHING_AT_LAST_LINE:

sep = b'' if i == (num_lines - 1) else sep_values[0]

elif eol == EOL.CUSTOM_DELIMITER:

if custom_delimiter is None or len(custom_delimiter) == 0:

raise ValueError('delimiter can not be null or empty')

else:

sep = custom_delimiter

else:

raise ValueError('Received unknown value %s for eol.' % eol)

f.write(data + sep)

return f.name, [line.decode('utf-8') for line in all_data]

def write_pattern(lines_per_file, no_data=False, return_filenames=False):

"""Writes a pattern of temporary files.

Args:

lines_per_file (List[int]): The number of lines to write per file.

no_data (bool): If :data:`True`, empty lines will be written, otherwise

each line will contain a concatenation of b'line' and the line number.

return_filenames (bool): If True, returned list will contain

(filename, data) pairs.

Returns:

Tuple[str, List[Union[str, (str, str)]]]: A tuple of the filename pattern

and a list of the utf-8 decoded written data or (filename, data) pairs.

"""

temp_dir = tempfile.mkdtemp()

all_data = []

file_name = None

start_index = 0

for i in range(len(lines_per_file)):

file_name, data = write_data(lines_per_file[i], no_data=no_data,

directory=temp_dir, prefix='mytemp')

if return_filenames:

all_data.extend(zip([file_name] * len(data), data))

else:

all_data.extend(data)

start_index += lines_per_file[i]

assert file_name

return (

file_name[:file_name.rfind(os.path.sep)] + os.path.sep + 'mytemp*',

all_data)

class TextSourceTest(unittest.TestCase):

# Number of records that will be written by most tests.

DEFAULT_NUM_RECORDS = 100

def _run_read_test(

self,

file_or_pattern,

expected_data,

buffer_size=DEFAULT_NUM_RECORDS,

compression=CompressionTypes.UNCOMPRESSED,

delimiter=None,

escapechar=None):

# Since each record usually takes more than 1 byte, default buffer size is

# smaller than the total size of the file. This is done to

# increase test coverage for cases that hit the buffer boundary.

kwargs = {}

if delimiter:

kwargs['delimiter'] = delimiter

if escapechar:

kwargs['escapechar'] = escapechar

source = TextSource(

file_or_pattern,

0,

compression,

True,

coders.StrUtf8Coder(),

buffer_size,

**kwargs)

range_tracker = source.get_range_tracker(None, None)

read_data = list(source.read(range_tracker))

self.assertCountEqual(expected_data, read_data)

@unittest.skipIf(platform.system() == 'Windows', 'Skipping on Windows')

def test_read_from_text_file_pattern_with_dot_slash(self):

cwd = os.getcwd()

expected = ['abc', 'de']

with TempDir() as temp_dir:

temp_dir.create_temp_file(suffix='.txt', lines=[b'a', b'b', b'c'])

temp_dir.create_temp_file(suffix='.txt', lines=[b'd', b'e'])

os.chdir(temp_dir.get_path())

with TestPipeline() as p:

dot_slash = p | 'ReadDotSlash' >> ReadFromText('./*.txt')

no_dot_slash = p | 'ReadNoSlash' >> ReadFromText('*.txt')

assert_that(dot_slash, equal_to(expected))

assert_that(no_dot_slash, equal_to(expected))

os.chdir(cwd)

def test_read_from_text_with_value_provider(self):

class UserDefinedOptions(PipelineOptions):

@classmethod

def _add_argparse_args(cls, parser):

parser.add_value_provider_argument(

'--file_pattern',

help='This keyword argument is a value provider',

default='some value')

options = UserDefinedOptions(['--file_pattern', 'abc'])

with self.assertRaises(OSError):

with TestPipeline(options=options) as pipeline:

_ = pipeline | 'Read' >> ReadFromText(options.file_pattern)

def test_read_single_file(self):

file_name, expected_data = write_data(TextSourceTest.DEFAULT_NUM_RECORDS)

assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS

self._run_read_test(file_name, expected_data)

def test_read_single_file_smaller_than_default_buffer(self):

file_name, expected_data = write_data(TextSourceTest.DEFAULT_NUM_RECORDS)

self._run_read_test(

file_name,

expected_data,

buffer_size=TextSource.DEFAULT_READ_BUFFER_SIZE)

def test_read_single_file_larger_than_default_buffer(self):

file_name, expected_data = write_data(TextSource.DEFAULT_READ_BUFFER_SIZE)

self._run_read_test(

file_name,

expected_data,

buffer_size=TextSource.DEFAULT_READ_BUFFER_SIZE)

def test_read_file_pattern(self):

pattern, expected_data = write_pattern(

[TextSourceTest.DEFAULT_NUM_RECORDS * 5,

TextSourceTest.DEFAULT_NUM_RECORDS * 3,

TextSourceTest.DEFAULT_NUM_RECORDS * 12,

TextSourceTest.DEFAULT_NUM_RECORDS * 8,

TextSourceTest.DEFAULT_NUM_RECORDS * 8,

TextSourceTest.DEFAULT_NUM_RECORDS * 4])

assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS * 40

self._run_read_test(pattern, expected_data)

def test_read_single_file_windows_eol(self):

file_name, expected_data = write_data(TextSourceTest.DEFAULT_NUM_RECORDS,

eol=EOL.CRLF)

assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS

self._run_read_test(file_name, expected_data)

def test_read_single_file_mixed_eol(self):

file_name, expected_data = write_data(TextSourceTest.DEFAULT_NUM_RECORDS,

eol=EOL.MIXED)

assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS

self._run_read_test(file_name, expected_data)

def test_read_single_file_last_line_no_eol(self):

file_name, expected_data = write_data(

TextSourceTest.DEFAULT_NUM_RECORDS,

eol=EOL.LF_WITH_NOTHING_AT_LAST_LINE)

assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS

self._run_read_test(file_name, expected_data)

def test_read_single_file_single_line_no_eol(self):

file_name, expected_data = write_data(

1, eol=EOL.LF_WITH_NOTHING_AT_LAST_LINE)

assert len(expected_data) == 1

self._run_read_test(file_name, expected_data)

def test_read_empty_single_file(self):

file_name, written_data = write_data(

1, no_data=True, eol=EOL.LF_WITH_NOTHING_AT_LAST_LINE)

assert len(written_data) == 1

# written data has a single entry with an empty string. Reading the source

# should not produce anything since we only wrote a single empty string

# without an end of line character.

self._run_read_test(file_name, [])

def test_read_single_file_last_line_no_eol_gzip(self):

file_name, expected_data = write_data(

TextSourceTest.DEFAULT_NUM_RECORDS,

eol=EOL.LF_WITH_NOTHING_AT_LAST_LINE)

gzip_file_name = file_name + '.gz'

with open(file_name, 'rb') as src, gzip.open(gzip_file_name, 'wb') as dst:

dst.writelines(src)

assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS

self._run_read_test(

gzip_file_name, expected_data, compression=CompressionTypes.GZIP)

def test_read_single_file_single_line_no_eol_gzip(self):

file_name, expected_data = write_data(

1, eol=EOL.LF_WITH_NOTHING_AT_LAST_LINE)

gzip_file_name = file_name + '.gz'

with open(file_name, 'rb') as src, gzip.open(gzip_file_name, 'wb') as dst:

dst.writelines(src)

assert len(expected_data) == 1

self._run_read_test(

gzip_file_name, expected_data, compression=CompressionTypes.GZIP)

def test_read_empty_single_file_no_eol_gzip(self):

file_name, written_data = write_data(

1, no_data=True, eol=EOL.LF_WITH_NOTHING_AT_LAST_LINE)

gzip_file_name = file_name + '.gz'

with open(file_name, 'rb') as src, gzip.open(gzip_file_name, 'wb') as dst:

dst.writelines(src)

assert len(written_data) == 1

# written data has a single entry with an empty string. Reading the source

# should not produce anything since we only wrote a single empty string

# without an end of line character.

self._run_read_test(gzip_file_name, [], compression=CompressionTypes.GZIP)

def test_read_single_file_with_empty_lines(self):

file_name, expected_data = write_data(

TextSourceTest.DEFAULT_NUM_RECORDS, no_data=True, eol=EOL.LF)

assert len(expected_data) == TextSourceTest.DEFAULT_NUM_RECORDS

assert not expected_data[0]

self._run_read_test(file_name, expected_data)

def test_read_single_file_without_striping_eol_lf(self):

file_name, written_data = write_data(TextSourceTest.DEFAULT_NUM_RECORDS,

eol=EOL.LF)

assert len(written_data) == TextSourceTest.DEFAULT_NUM_RECORDS

source = TextSource(

file_name,

0,

CompressionTypes.UNCOMPRESSED,

False,

coders.StrUtf8Coder())

range_tracker = source.get_range_tracker(None, None)

read_data = list(source.read(range_tracker))

self.assertCountEqual([line + '\n' for line in written_data], read_data)

def test_read_single_file_without_striping_eol_crlf(self):

file_name, written_data = write_data(TextSourceTest.DEFAULT_NUM_RECORDS,

eol=EOL.CRLF)

assert len(written_data) == TextSourceTest.DEFAULT_NUM_RECORDS

source = TextSource(

file_name,

0,

CompressionTypes.UNCOMPRESSED,

False,

coders.StrUtf8Coder())

range_tracker = source.get_range_tracker(None, None)

read_data = list(source.read(range_tracker))

self.assertCountEqual([line + '\r\n' for line in written_data], read_data)

def test_read_file_pattern_with_empty_files(self):

pattern, expected_data = write_pattern(

[5 * TextSourceTest.DEFAULT_NUM_RECORDS,

3 * TextSourceTest.DEFAULT_NUM_RECORDS,

12 * TextSourceTest.DEFAULT_NUM_RECORDS,

8 * TextSourceTest.DEFAULT_NUM_RECORDS,

8 * TextSourceTest.DEFAULT_NUM_RECORDS,

4 * TextSourceTest.DEFAULT_NUM_RECORDS],

no_data=True)

assert len(expected_data) == 40 * TextSourceTest.DEFAULT_NUM_RECORDS

assert not expected_data[0]

self._run_read_test(pattern, expected_data)

def test_read_after_splitting(self):

file_name, expected_data = write_data(10)

assert len(expected_data) == 10

source = TextSource(

file_name,

0,

CompressionTypes.UNCOMPRESSED,

True,

coders.StrUtf8Coder())

splits = list(source.split(desired_bundle_size=33))

reference_source_info = (source, None, None)

sources_info = ([(split.source, split.start_position, split.stop_position)

for split in splits])

source_test_utils.assert_sources_equal_reference_source(

reference_source_info, sources_info)

def test_header_processing(self):

file_name, expected_data = write_data(10)

assert len(expected_data) == 10

def header_matcher(line):

return line in expected_data[:5]

header_lines = []

def store_header(lines):

for line in lines:

header_lines.append(line)

source = TextSource(

file_name,

0,

CompressionTypes.UNCOMPRESSED,

True,

coders.StrUtf8Coder(),

header_processor_fns=(header_matcher, store_header))

splits = list(source.split(desired_bundle_size=100000))

assert len(splits) == 1

range_tracker = splits[0].source.get_range_tracker(

splits[0].start_position, splits[0].stop_position)

read_data = list(source.read_records(file_name, range_tracker))

self.assertCountEqual(expected_data[:5], header_lines)

self.assertCountEqual(expected_data[5:], read_data)

def test_progress(self):

file_name, expected_data = write_data(10)

assert len(expected_data) == 10

source = TextSource(

file_name,

0,

CompressionTypes.UNCOMPRESSED,

True,

coders.StrUtf8Coder())

splits = list(source.split(desired_bundle_size=100000))

assert len(splits) == 1

fraction_consumed_report = []

split_points_report = []

range_tracker = splits[0].source.get_range_tracker(

splits[0].start_position, splits[0].stop_position)

for _ in splits[0].source.read(range_tracker):

fraction_consumed_report.append(range_tracker.fraction_consumed())

split_points_report.append(range_tracker.split_points())

self.assertEqual([float(i) / 10 for i in range(0, 10)],

fraction_consumed_report)

expected_split_points_report = [((i - 1),

iobase.RangeTracker.SPLIT_POINTS_UNKNOWN)

for i in range(1, 10)]

# At last split point, the remaining split points callback returns 1 since

# the expected position of next record becomes equal to the stop position.

expected_split_points_report.append((9, 1))

self.assertEqual(expected_split_points_report, split_points_report)

def test_read_reentrant_without_splitting(self):

file_name, expected_data = write_data(10)

assert len(expected_data) == 10

source = TextSource(

file_name,

0,

CompressionTypes.UNCOMPRESSED,

True,

coders.StrUtf8Coder())

source_test_utils.assert_reentrant_reads_succeed((source, None, None))

def test_read_reentrant_after_splitting(self):

file_name, expected_data = write_data(10)

assert len(expected_data) == 10

source = TextSource(

file_name,

0,

CompressionTypes.UNCOMPRESSED,

True,

coders.StrUtf8Coder())

splits = list(source.split(desired_bundle_size=100000))

assert len(splits) == 1

source_test_utils.assert_reentrant_reads_succeed(

(splits[0].source, splits[0].start_position, splits[0].stop_position))

def test_dynamic_work_rebalancing(self):

file_name, expected_data = write_data(5)

assert len(expected_data) == 5

source = TextSource(

file_name,

0,

CompressionTypes.UNCOMPRESSED,

True,

coders.StrUtf8Coder())

splits = list(source.split(desired_bundle_size=100000))

assert len(splits) == 1

source_test_utils.assert_split_at_fraction_exhaustive(

splits[0].source, splits[0].start_position, splits[0].stop_position)

def test_dynamic_work_rebalancing_windows_eol(self):

file_name, expected_data = write_data(15, eol=EOL.CRLF)

assert len(expected_data) == 15

source = TextSource(

file_name,

0,

CompressionTypes.UNCOMPRESSED,

True,

coders.StrUtf8Coder())

splits = list(source.split(desired_bundle_size=100000))

assert len(splits) == 1

source_test_utils.assert_split_at_fraction_exhaustive(

splits[0].source,

splits[0].start_position,

splits[0].stop_position,

perform_multi_threaded_test=False)

def test_dynamic_work_rebalancing_mixed_eol(self):

file_name, expected_data = write_data(5, eol=EOL.MIXED)

assert len(expected_data) == 5

source = TextSource(

file_name,

0,

CompressionTypes.UNCOMPRESSED,

True,

coders.StrUtf8Coder())

splits = list(source.split(desired_bundle_size=100000))

assert len(splits) == 1

source_test_utils.assert_split_at_fraction_exhaustive(

splits[0].source,

splits[0].start_position,

splits[0].stop_position,

perform_multi_threaded_test=False)

def test_read_from_text_single_file(self):

file_name, expected_data = write_data(5)

assert len(expected_data) == 5

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromText(file_name)

assert_that(pcoll, equal_to(expected_data))

def test_read_from_text_with_file_name_single_file(self):

file_name, data = write_data(5)

expected_data = [(file_name, el) for el in data]

assert len(expected_data) == 5

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromTextWithFilename(file_name)

assert_that(pcoll, equal_to(expected_data))

def test_read_all_single_file(self):

file_name, expected_data = write_data(5)

assert len(expected_data) == 5

with TestPipeline() as pipeline:

pcoll = pipeline | 'Create' >> Create(

[file_name]) | 'ReadAll' >> ReadAllFromText()

assert_that(pcoll, equal_to(expected_data))

def test_read_all_many_single_files(self):

file_name1, expected_data1 = write_data(5)

assert len(expected_data1) == 5

file_name2, expected_data2 = write_data(10)

assert len(expected_data2) == 10

file_name3, expected_data3 = write_data(15)

assert len(expected_data3) == 15

expected_data = []

expected_data.extend(expected_data1)

expected_data.extend(expected_data2)

expected_data.extend(expected_data3)

with TestPipeline() as pipeline:

pcoll = pipeline | 'Create' >> Create([

file_name1, file_name2, file_name3

]) | 'ReadAll' >> ReadAllFromText()

assert_that(pcoll, equal_to(expected_data))

def test_read_all_unavailable_files_ignored(self):

file_name1, expected_data1 = write_data(5)

assert len(expected_data1) == 5

file_name2, expected_data2 = write_data(10)

assert len(expected_data2) == 10

file_name3, expected_data3 = write_data(15)

assert len(expected_data3) == 15

file_name4 = "/unavailable_file"

expected_data = []

expected_data.extend(expected_data1)

expected_data.extend(expected_data2)

expected_data.extend(expected_data3)

with TestPipeline() as pipeline:

pcoll = (

pipeline

| 'Create' >> Create([file_name1, file_name2, file_name3, file_name4])

| 'ReadAll' >> ReadAllFromText())

assert_that(pcoll, equal_to(expected_data))

class _WriteFilesFn(beam.DoFn):

"""writes a couple of files with deferral."""

COUNT_STATE = CombiningValueStateSpec('count', combine_fn=sum)

def __init__(self, temp_path):

self.temp_path = temp_path

def process(self, element, count_state=beam.DoFn.StateParam(COUNT_STATE)):

counter = count_state.read()

if counter == 0:

count_state.add(1)

with open(FileSystems.join(self.temp_path, 'file1'), 'w') as f:

f.write('second A\nsecond B')

with open(FileSystems.join(self.temp_path, 'file2'), 'w') as f:

f.write('first')

# convert dumb key to basename in output

basename = FileSystems.split(element[1][0])[1]

content = element[1][1]

yield basename, content

def test_read_all_continuously_new(self):

with TempDir() as tempdir, TestPipeline() as pipeline:

temp_path = tempdir.get_path()

# create a temp file at the beginning

with open(FileSystems.join(temp_path, 'file1'), 'w') as f:

f.write('first')

match_pattern = FileSystems.join(temp_path, '*')

interval = 0.5

last = 2

p_read_once = (

pipeline

| 'Continuously read new files' >> ReadAllFromTextContinuously(

match_pattern,

with_filename=True,

start_timestamp=Timestamp.now(),

interval=interval,

stop_timestamp=Timestamp.now() + last,

match_updated_files=False)

| 'add dumb key' >> beam.Map(lambda x: (0, x))

|

'Write files on-the-fly' >> beam.ParDo(self._WriteFilesFn(temp_path)))

assert_that(

p_read_once,

equal_to([('file1', 'first'), ('file2', 'first')]),

label='assert read new files results')

def test_read_all_continuously_update(self):

with TempDir() as tempdir, TestPipeline() as pipeline:

temp_path = tempdir.get_path()

# create a temp file at the beginning

with open(FileSystems.join(temp_path, 'file1'), 'w') as f:

f.write('first')

match_pattern = FileSystems.join(temp_path, '*')

interval = 0.5

last = 2

p_read_upd = (

pipeline

| 'Continuously read updated files' >> ReadAllFromTextContinuously(

match_pattern,

with_filename=True,

start_timestamp=Timestamp.now(),

interval=interval,

stop_timestamp=Timestamp.now() + last,

match_updated_files=True)

| 'add dumb key' >> beam.Map(lambda x: (0, x))

|

'Write files on-the-fly' >> beam.ParDo(self._WriteFilesFn(temp_path)))

assert_that(

p_read_upd,

equal_to([('file1', 'first'), ('file1', 'second A'),

('file1', 'second B'), ('file2', 'first')]),

label='assert read updated files results')

def test_read_from_text_single_file_with_coder(self):

file_name, expected_data = write_data(5)

assert len(expected_data) == 5

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromText(file_name, coder=DummyCoder())

assert_that(pcoll, equal_to([record * 2 for record in expected_data]))

def test_read_from_text_file_pattern(self):

pattern, expected_data = write_pattern([5, 3, 12, 8, 8, 4])

assert len(expected_data) == 40

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromText(pattern)

assert_that(pcoll, equal_to(expected_data))

def test_read_from_text_with_file_name_file_pattern(self):

pattern, expected_data = write_pattern(

lines_per_file=[5, 5], return_filenames=True)

assert len(expected_data) == 10

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromTextWithFilename(pattern)

assert_that(pcoll, equal_to(expected_data))

def test_read_all_file_pattern(self):

pattern, expected_data = write_pattern([5, 3, 12, 8, 8, 4])

assert len(expected_data) == 40

with TestPipeline() as pipeline:

pcoll = (

pipeline

| 'Create' >> Create([pattern])

| 'ReadAll' >> ReadAllFromText())

assert_that(pcoll, equal_to(expected_data))

def test_read_all_many_file_patterns(self):

pattern1, expected_data1 = write_pattern([5, 3, 12, 8, 8, 4])

assert len(expected_data1) == 40

pattern2, expected_data2 = write_pattern([3, 7, 9])

assert len(expected_data2) == 19

pattern3, expected_data3 = write_pattern([11, 20, 5, 5])

assert len(expected_data3) == 41

expected_data = []

expected_data.extend(expected_data1)

expected_data.extend(expected_data2)

expected_data.extend(expected_data3)

with TestPipeline() as pipeline:

pcoll = pipeline | 'Create' >> Create(

[pattern1, pattern2, pattern3]) | 'ReadAll' >> ReadAllFromText()

assert_that(pcoll, equal_to(expected_data))

def test_read_all_with_filename(self):

pattern, expected_data = write_pattern([5, 3], return_filenames=True)

assert len(expected_data) == 8

with TestPipeline() as pipeline:

pcoll = (

pipeline

| 'Create' >> Create([pattern])

| 'ReadAll' >> ReadAllFromText(with_filename=True))

assert_that(pcoll, equal_to(expected_data))

def test_read_auto_bzip2(self):

_, lines = write_data(15)

with TempDir() as tempdir:

file_name = tempdir.create_temp_file(suffix='.bz2')

with bz2.BZ2File(file_name, 'wb') as f:

f.write('\n'.join(lines).encode('utf-8'))

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromText(file_name)

assert_that(pcoll, equal_to(lines))

def test_read_auto_deflate(self):

_, lines = write_data(15)

with TempDir() as tempdir:

file_name = tempdir.create_temp_file(suffix='.deflate')

with open(file_name, 'wb') as f:

f.write(zlib.compress('\n'.join(lines).encode('utf-8')))

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromText(file_name)

assert_that(pcoll, equal_to(lines))

def test_read_auto_gzip(self):

_, lines = write_data(15)

with TempDir() as tempdir:

file_name = tempdir.create_temp_file(suffix='.gz')

with gzip.GzipFile(file_name, 'wb') as f:

f.write('\n'.join(lines).encode('utf-8'))

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromText(file_name)

assert_that(pcoll, equal_to(lines))

def test_read_bzip2(self):

_, lines = write_data(15)

with TempDir() as tempdir:

file_name = tempdir.create_temp_file()

with bz2.BZ2File(file_name, 'wb') as f:

f.write('\n'.join(lines).encode('utf-8'))

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromText(

file_name, compression_type=CompressionTypes.BZIP2)

assert_that(pcoll, equal_to(lines))

def test_read_corrupted_bzip2_fails(self):

_, lines = write_data(15)

with TempDir() as tempdir:

file_name = tempdir.create_temp_file()

with bz2.BZ2File(file_name, 'wb') as f:

f.write('\n'.join(lines).encode('utf-8'))

with open(file_name, 'wb') as f:

f.write(b'corrupt')

with self.assertRaises(Exception):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromText(

file_name, compression_type=CompressionTypes.BZIP2)

assert_that(pcoll, equal_to(lines))

def test_read_bzip2_concat(self):

with TempDir() as tempdir:

bzip2_file_name1 = tempdir.create_temp_file()

lines = ['a', 'b', 'c']

with bz2.BZ2File(bzip2_file_name1, 'wb') as dst:

data = '\n'.join(lines) + '\n'

dst.write(data.encode('utf-8'))

bzip2_file_name2 = tempdir.create_temp_file()

lines = ['p', 'q', 'r']

with bz2.BZ2File(bzip2_file_name2, 'wb') as dst:

data = '\n'.join(lines) + '\n'

dst.write(data.encode('utf-8'))

bzip2_file_name3 = tempdir.create_temp_file()

lines = ['x', 'y', 'z']

with bz2.BZ2File(bzip2_file_name3, 'wb') as dst:

data = '\n'.join(lines) + '\n'

dst.write(data.encode('utf-8'))

final_bzip2_file = tempdir.create_temp_file()

with open(bzip2_file_name1, 'rb') as src, open(

final_bzip2_file, 'wb') as dst:

dst.writelines(src.readlines())

with open(bzip2_file_name2, 'rb') as src, open(

final_bzip2_file, 'ab') as dst:

dst.writelines(src.readlines())

with open(bzip2_file_name3, 'rb') as src, open(

final_bzip2_file, 'ab') as dst:

dst.writelines(src.readlines())

with TestPipeline() as pipeline:

lines = pipeline | 'ReadFromText' >> beam.io.ReadFromText(

final_bzip2_file,

compression_type=beam.io.filesystem.CompressionTypes.BZIP2)

expected = ['a', 'b', 'c', 'p', 'q', 'r', 'x', 'y', 'z']

assert_that(lines, equal_to(expected))

def test_read_deflate(self):

_, lines = write_data(15)

with TempDir() as tempdir:

file_name = tempdir.create_temp_file()

with open(file_name, 'wb') as f:

f.write(zlib.compress('\n'.join(lines).encode('utf-8')))

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromText(

file_name, 0, CompressionTypes.DEFLATE, True, coders.StrUtf8Coder())

assert_that(pcoll, equal_to(lines))

def test_read_corrupted_deflate_fails(self):

_, lines = write_data(15)

with TempDir() as tempdir:

file_name = tempdir.create_temp_file()

with open(file_name, 'wb') as f:

f.write(zlib.compress('\n'.join(lines).encode('utf-8')))

with open(file_name, 'wb') as f:

f.write(b'corrupt')

with self.assertRaises(Exception):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromText(

file_name,

0,

CompressionTypes.DEFLATE,

True,

coders.StrUtf8Coder())

assert_that(pcoll, equal_to(lines))

def test_read_deflate_concat(self):

with TempDir() as tempdir:

deflate_file_name1 = tempdir.create_temp_file()

lines = ['a', 'b', 'c']

with open(deflate_file_name1, 'wb') as dst:

data = '\n'.join(lines) + '\n'

dst.write(zlib.compress(data.encode('utf-8')))

deflate_file_name2 = tempdir.create_temp_file()

lines = ['p', 'q', 'r']

with open(deflate_file_name2, 'wb') as dst:

data = '\n'.join(lines) + '\n'

dst.write(zlib.compress(data.encode('utf-8')))

deflate_file_name3 = tempdir.create_temp_file()

lines = ['x', 'y', 'z']

with open(deflate_file_name3, 'wb') as dst:

data = '\n'.join(lines) + '\n'

dst.write(zlib.compress(data.encode('utf-8')))

final_deflate_file = tempdir.create_temp_file()

with open(deflate_file_name1, 'rb') as src, \

open(final_deflate_file, 'wb') as dst:

dst.writelines(src.readlines())

with open(deflate_file_name2, 'rb') as src, \

open(final_deflate_file, 'ab') as dst:

dst.writelines(src.readlines())

with open(deflate_file_name3, 'rb') as src, \

open(final_deflate_file, 'ab') as dst:

dst.writelines(src.readlines())

with TestPipeline() as pipeline:

lines = pipeline | 'ReadFromText' >> beam.io.ReadFromText(

final_deflate_file,

compression_type=beam.io.filesystem.CompressionTypes.DEFLATE)

expected = ['a', 'b', 'c', 'p', 'q', 'r', 'x', 'y', 'z']

assert_that(lines, equal_to(expected))

def test_read_gzip(self):

_, lines = write_data(15)

with TempDir() as tempdir:

file_name = tempdir.create_temp_file()

with gzip.GzipFile(file_name, 'wb') as f:

f.write('\n'.join(lines).encode('utf-8'))

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromText(

file_name, 0, CompressionTypes.GZIP, True, coders.StrUtf8Coder())

assert_that(pcoll, equal_to(lines))

def test_read_corrupted_gzip_fails(self):

_, lines = write_data(15)

with TempDir() as tempdir:

file_name = tempdir.create_temp_file()

with gzip.GzipFile(file_name, 'wb') as f:

f.write('\n'.join(lines).encode('utf-8'))

with open(file_name, 'wb') as f:

f.write(b'corrupt')

with self.assertRaises(Exception):

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromText(

file_name, 0, CompressionTypes.GZIP, True, coders.StrUtf8Coder())

assert_that(pcoll, equal_to(lines))

def test_read_gzip_concat(self):

with TempDir() as tempdir:

gzip_file_name1 = tempdir.create_temp_file()

lines = ['a', 'b', 'c']

with gzip.open(gzip_file_name1, 'wb') as dst:

data = '\n'.join(lines) + '\n'

dst.write(data.encode('utf-8'))

gzip_file_name2 = tempdir.create_temp_file()

lines = ['p', 'q', 'r']

with gzip.open(gzip_file_name2, 'wb') as dst:

data = '\n'.join(lines) + '\n'

dst.write(data.encode('utf-8'))

gzip_file_name3 = tempdir.create_temp_file()

lines = ['x', 'y', 'z']

with gzip.open(gzip_file_name3, 'wb') as dst:

data = '\n'.join(lines) + '\n'

dst.write(data.encode('utf-8'))

final_gzip_file = tempdir.create_temp_file()

with open(gzip_file_name1, 'rb') as src, \

open(final_gzip_file, 'wb') as dst:

dst.writelines(src.readlines())

with open(gzip_file_name2, 'rb') as src, \

open(final_gzip_file, 'ab') as dst:

dst.writelines(src.readlines())

with open(gzip_file_name3, 'rb') as src, \

open(final_gzip_file, 'ab') as dst:

dst.writelines(src.readlines())

with TestPipeline() as pipeline:

lines = pipeline | 'ReadFromText' >> beam.io.ReadFromText(

final_gzip_file,

compression_type=beam.io.filesystem.CompressionTypes.GZIP)

expected = ['a', 'b', 'c', 'p', 'q', 'r', 'x', 'y', 'z']

assert_that(lines, equal_to(expected))

def test_read_all_gzip(self):

_, lines = write_data(100)

with TempDir() as tempdir:

file_name = tempdir.create_temp_file()

with gzip.GzipFile(file_name, 'wb') as f:

f.write('\n'.join(lines).encode('utf-8'))

with TestPipeline() as pipeline:

pcoll = (

pipeline

| Create([file_name])

| 'ReadAll' >>

ReadAllFromText(compression_type=CompressionTypes.GZIP))

assert_that(pcoll, equal_to(lines))

def test_read_gzip_large(self):

_, lines = write_data(10000)

with TempDir() as tempdir:

file_name = tempdir.create_temp_file()

with gzip.GzipFile(file_name, 'wb') as f:

f.write('\n'.join(lines).encode('utf-8'))

with TestPipeline() as pipeline:

pcoll = pipeline | 'Read' >> ReadFromText(