<s> """<STR_LIT>""" <EOL> import base64 <EOL> import logging <EOL> import unittest <EOL> from google . cloud . dataflow import coders <EOL> from google . cloud . dataflow import io <EOL> from google . cloud . dataflow . internal . json_value import to_json_value <EOL> from google . cloud . dataflow . worker import concat_reader <EOL> from google . cloud . dataflow . worker import inmemory <EOL> from google . cloud . dataflow . worker import maptask <EOL> from google . cloud . dataflow . worker import workitem <EOL> import google . cloud . dataflow . internal . clients . dataflow as dataflow <EOL> TEXT_SOURCE_SPEC = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT:filename>' : { '<STR_LIT:value>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : True , '<STR_LIT>' : '<STR_LIT>' } , <EOL> } <EOL> IN_MEMORY_ELEMENTS = [ <EOL> { '<STR_LIT:value>' : base64 . b64encode ( '<STR_LIT:1>' ) , '<STR_LIT>' : '<STR_LIT>' } , <EOL> { '<STR_LIT:value>' : base64 . b64encode ( '<STR_LIT:2>' ) , '<STR_LIT>' : '<STR_LIT>' } , <EOL> { '<STR_LIT:value>' : base64 . b64encode ( '<STR_LIT:3>' ) , '<STR_LIT>' : '<STR_LIT>' } ] <EOL> IN_MEMORY_SOURCE_SPEC = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT:1>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT:3>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : IN_MEMORY_ELEMENTS , <EOL> } <EOL> GROUPING_SHUFFLE_SOURCE_SPEC = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> } <EOL> UNGROUPED_SHUFFLE_SOURCE_SPEC = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> } <EOL> TEXT_SINK_SPEC = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT:filename>' : { '<STR_LIT:value>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : True , '<STR_LIT>' : '<STR_LIT>' } , <EOL> } <EOL> SHUFFLE_SINK_SPEC = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> } <EOL> PARDO_DOFN_SPEC = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT:code>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> } <EOL> PARDO_COMBINEFN_SPEC = { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT:code>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : { '<STR_LIT:value>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } , <EOL> } <EOL> CONCAT_SOURCE_SPEC = { <EOL> '<STR_LIT>' : <EOL> [ <EOL> { <EOL> '<STR_LIT>' : { <EOL> '<STR_LIT>' : <EOL> { '<STR_LIT>' : '<STR_LIT>' , '<STR_LIT:value>' : True } , <EOL> '<STR_LIT>' : <EOL> { '<STR_LIT>' : '<STR_LIT>' , '<STR_LIT:value>' : '<STR_LIT:0>' } , <EOL> '<STR_LIT>' : <EOL> { '<STR_LIT>' : '<STR_LIT>' , '<STR_LIT:value>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : <EOL> { '<STR_LIT>' : '<STR_LIT>' , '<STR_LIT:value>' : '<STR_LIT>' } , <EOL> '<STR_LIT:filename>' : <EOL> { '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT:value>' : '<STR_LIT>' <EOL> '<STR_LIT>' } , <EOL> '<STR_LIT>' : '<STR_LIT>' <EOL> } , <EOL> '<STR_LIT>' : { <EOL> '<STR_LIT>' : <EOL> [ { '<STR_LIT>' : '<STR_LIT>' } , { '<STR_LIT>' : '<STR_LIT>' } ] , <EOL> '<STR_LIT>' : coders . serialize_coder ( coders . PickleCoder ( ) ) <EOL> } <EOL> } , <EOL> { <EOL> '<STR_LIT>' : { <EOL> '<STR_LIT>' : <EOL> { '<STR_LIT>' : '<STR_LIT>' , '<STR_LIT:value>' : True } , <EOL> '<STR_LIT>' : <EOL> { '<STR_LIT>' : '<STR_LIT>' , '<STR_LIT:value>' : '<STR_LIT:0>' } , <EOL> '<STR_LIT>' : <EOL> { '<STR_LIT>' : '<STR_LIT>' , '<STR_LIT:value>' : '<STR_LIT>' } , <EOL> '<STR_LIT>' : <EOL> { '<STR_LIT>' : '<STR_LIT>' , '<STR_LIT:value>' : '<STR_LIT>' } , <EOL> '<STR_LIT:filename>' : <EOL> { '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT:value>' : '<STR_LIT>' <EOL> '<STR_LIT>' } , <EOL> '<STR_LIT>' : '<STR_LIT>' <EOL> } , <EOL> '<STR_LIT>' : { <EOL> '<STR_LIT>' : <EOL> [ { '<STR_LIT>' : '<STR_LIT>' } , { '<STR_LIT>' : '<STR_LIT>' } ] , <EOL> '<STR_LIT>' : coders . serialize_coder ( coders . PickleCoder ( ) ) <EOL> } <EOL> } <EOL> ] , <EOL> '<STR_LIT>' : '<STR_LIT>' <EOL> } <EOL> CODER = coders . PickleCoder ( ) <EOL> WINDOWED_CODER = coders . WindowedValueCoder ( CODER ) <EOL> CODER_SPEC = CODER . as_cloud_object ( ) <EOL> WINDOWED_CODER_SPEC = WINDOWED_CODER . as_cloud_object ( ) <EOL> def add_source_codec_spec ( target ) : <EOL> target . source . codec = dataflow . Source . CodecValue ( ) <EOL> for k , v in CODER_SPEC . iteritems ( ) : <EOL> target . source . codec . additionalProperties . append ( <EOL> dataflow . Source . CodecValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> def add_source_windowed_codec_spec ( target ) : <EOL> target . source . codec = dataflow . Source . CodecValue ( ) <EOL> for k , v in WINDOWED_CODER_SPEC . iteritems ( ) : <EOL> target . source . codec . additionalProperties . append ( <EOL> dataflow . Source . CodecValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> def add_sink_codec_spec ( target ) : <EOL> target . sink . codec = dataflow . Sink . CodecValue ( ) <EOL> for k , v in CODER_SPEC . iteritems ( ) : <EOL> target . sink . codec . additionalProperties . append ( <EOL> dataflow . Sink . CodecValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> def get_instruction_with_outputs ( num_outputs = <NUM_LIT:1> , ** kwargs ) : <EOL> pi = dataflow . ParallelInstruction ( ** kwargs ) <EOL> for _ in xrange ( num_outputs ) : <EOL> output = dataflow . InstructionOutput ( ) <EOL> output . codec = dataflow . InstructionOutput . CodecValue ( ) <EOL> for k , v in CODER_SPEC . iteritems ( ) : <EOL> output . codec . additionalProperties . append ( <EOL> dataflow . InstructionOutput . CodecValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> pi . outputs . append ( output ) <EOL> return pi <EOL> def get_concat_source_to_shuffle_sink_message ( ) : <EOL> ri = dataflow . ReadInstruction ( ) <EOL> ri . source = dataflow . Source ( ) <EOL> ri . source . spec = dataflow . Source . SpecValue ( ) <EOL> for k , v in CONCAT_SOURCE_SPEC . iteritems ( ) : <EOL> ri . source . spec . additionalProperties . append ( <EOL> dataflow . Source . SpecValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> di = dataflow . ParDoInstruction ( ) <EOL> di . input = dataflow . InstructionInput ( ) <EOL> di . input . producerInstructionIndex = <NUM_LIT:1> <EOL> di . multiOutputInfos = [ dataflow . MultiOutputInfo ( tag = '<STR_LIT>' ) ] <EOL> di . userFn = dataflow . ParDoInstruction . UserFnValue ( ) <EOL> for k , v in PARDO_DOFN_SPEC . iteritems ( ) : <EOL> di . userFn . additionalProperties . append ( <EOL> dataflow . ParDoInstruction . UserFnValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> wsi = dataflow . WriteInstruction ( ) <EOL> wsi . input = dataflow . InstructionInput ( ) <EOL> wsi . input . producerInstructionIndex = <NUM_LIT:1> <EOL> di . input . outputNum = <NUM_LIT:0> <EOL> wsi . sink = dataflow . Sink ( ) <EOL> wsi . sink . spec = dataflow . Sink . SpecValue ( ) <EOL> for k , v in SHUFFLE_SINK_SPEC . iteritems ( ) : <EOL> wsi . sink . spec . additionalProperties . append ( <EOL> dataflow . Sink . SpecValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> add_sink_codec_spec ( wsi ) <EOL> mt = dataflow . MapTask ( ) <EOL> mt . instructions . append ( get_instruction_with_outputs ( read = ri ) ) <EOL> mt . instructions . append ( get_instruction_with_outputs ( parDo = di ) ) <EOL> mt . instructions . append ( dataflow . ParallelInstruction ( write = wsi ) ) <EOL> wi = dataflow . WorkItem ( ) <EOL> wi . id = <NUM_LIT> <EOL> wi . projectId = '<STR_LIT>' <EOL> wi . jobId = '<STR_LIT>' <EOL> wi . mapTask = mt <EOL> m = dataflow . LeaseWorkItemResponse ( ) <EOL> m . workItems . append ( wi ) <EOL> return m <EOL> def get_text_source_to_shuffle_sink_message ( ) : <EOL> ri = dataflow . ReadInstruction ( ) <EOL> ri . source = dataflow . Source ( ) <EOL> ri . source . spec = dataflow . Source . SpecValue ( ) <EOL> for k , v in TEXT_SOURCE_SPEC . iteritems ( ) : <EOL> ri . source . spec . additionalProperties . append ( <EOL> dataflow . Source . SpecValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> add_source_codec_spec ( ri ) <EOL> di = dataflow . ParDoInstruction ( ) <EOL> di . input = dataflow . InstructionInput ( ) <EOL> di . input . producerInstructionIndex = <NUM_LIT:1> <EOL> di . multiOutputInfos = [ dataflow . MultiOutputInfo ( tag = '<STR_LIT>' ) ] <EOL> di . userFn = dataflow . ParDoInstruction . UserFnValue ( ) <EOL> for k , v in PARDO_DOFN_SPEC . iteritems ( ) : <EOL> di . userFn . additionalProperties . append ( <EOL> dataflow . ParDoInstruction . UserFnValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> wsi = dataflow . WriteInstruction ( ) <EOL> wsi . input = dataflow . InstructionInput ( ) <EOL> wsi . input . producerInstructionIndex = <NUM_LIT:1> <EOL> di . input . outputNum = <NUM_LIT:0> <EOL> wsi . sink = dataflow . Sink ( ) <EOL> wsi . sink . spec = dataflow . Sink . SpecValue ( ) <EOL> for k , v in SHUFFLE_SINK_SPEC . iteritems ( ) : <EOL> wsi . sink . spec . additionalProperties . append ( <EOL> dataflow . Sink . SpecValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> add_sink_codec_spec ( wsi ) <EOL> mt = dataflow . MapTask ( ) <EOL> mt . instructions . append ( get_instruction_with_outputs ( read = ri ) ) <EOL> mt . instructions . append ( get_instruction_with_outputs ( parDo = di ) ) <EOL> mt . instructions . append ( dataflow . ParallelInstruction ( write = wsi ) ) <EOL> wi = dataflow . WorkItem ( ) <EOL> wi . id = <NUM_LIT> <EOL> wi . projectId = '<STR_LIT>' <EOL> wi . jobId = '<STR_LIT>' <EOL> wi . mapTask = mt <EOL> m = dataflow . LeaseWorkItemResponse ( ) <EOL> m . workItems . append ( wi ) <EOL> return m <EOL> def get_shuffle_source_to_text_sink_message ( shuffle_source_spec ) : <EOL> rsi = dataflow . ReadInstruction ( ) <EOL> rsi . source = dataflow . Source ( ) <EOL> rsi . source . spec = dataflow . Source . SpecValue ( ) <EOL> for k , v in shuffle_source_spec . iteritems ( ) : <EOL> rsi . source . spec . additionalProperties . append ( <EOL> dataflow . Source . SpecValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> add_source_codec_spec ( rsi ) <EOL> wi = dataflow . WriteInstruction ( ) <EOL> wi . input = dataflow . InstructionInput ( ) <EOL> wi . sink = dataflow . Sink ( ) <EOL> wi . sink . spec = dataflow . Sink . SpecValue ( ) <EOL> for k , v in TEXT_SINK_SPEC . iteritems ( ) : <EOL> wi . sink . spec . additionalProperties . append ( <EOL> dataflow . Sink . SpecValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> add_sink_codec_spec ( wi ) <EOL> mt = dataflow . MapTask ( ) <EOL> mt . instructions . append ( get_instruction_with_outputs ( read = rsi ) ) <EOL> mt . instructions . append ( dataflow . ParallelInstruction ( write = wi ) ) <EOL> wi = dataflow . WorkItem ( ) <EOL> wi . id = <NUM_LIT> <EOL> wi . projectId = '<STR_LIT>' <EOL> wi . jobId = '<STR_LIT>' <EOL> wi . mapTask = mt <EOL> m = dataflow . LeaseWorkItemResponse ( ) <EOL> m . workItems . append ( wi ) <EOL> return m <EOL> def get_in_memory_source_to_text_sink_message ( ) : <EOL> rsi = dataflow . ReadInstruction ( ) <EOL> rsi . source = dataflow . Source ( ) <EOL> rsi . source . spec = dataflow . Source . SpecValue ( ) <EOL> for k , v in IN_MEMORY_SOURCE_SPEC . iteritems ( ) : <EOL> rsi . source . spec . additionalProperties . append ( <EOL> dataflow . Source . SpecValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> add_source_windowed_codec_spec ( rsi ) <EOL> wi = dataflow . WriteInstruction ( ) <EOL> wi . input = dataflow . InstructionInput ( ) <EOL> wi . sink = dataflow . Sink ( ) <EOL> wi . sink . spec = dataflow . Sink . SpecValue ( ) <EOL> for k , v in TEXT_SINK_SPEC . iteritems ( ) : <EOL> wi . sink . spec . additionalProperties . append ( <EOL> dataflow . Sink . SpecValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> add_sink_codec_spec ( wi ) <EOL> mt = dataflow . MapTask ( ) <EOL> mt . instructions . append ( get_instruction_with_outputs ( read = rsi ) ) <EOL> mt . instructions . append ( dataflow . ParallelInstruction ( write = wi ) ) <EOL> wi = dataflow . WorkItem ( ) <EOL> wi . id = <NUM_LIT> <EOL> wi . projectId = '<STR_LIT>' <EOL> wi . jobId = '<STR_LIT>' <EOL> wi . mapTask = mt <EOL> m = dataflow . LeaseWorkItemResponse ( ) <EOL> m . workItems . append ( wi ) <EOL> return m <EOL> def get_in_memory_source_to_flatten_message ( ) : <EOL> rsi = dataflow . ReadInstruction ( ) <EOL> rsi . source = dataflow . Source ( ) <EOL> add_source_codec_spec ( rsi ) <EOL> rsi . source . spec = dataflow . Source . SpecValue ( ) <EOL> for k , v in IN_MEMORY_SOURCE_SPEC . iteritems ( ) : <EOL> rsi . source . spec . additionalProperties . append ( <EOL> dataflow . Source . SpecValue . AdditionalProperty ( <EOL> key = k , value = to_json_value ( v ) ) ) <EOL> add_source_windowed_codec_spec ( rsi ) <EOL> fi = dataflow . FlattenInstruction ( ) <EOL> fi . inputs = [ dataflow . InstructionInput ( ) ] <EOL> mt = dataflow . MapTask ( ) <EOL> mt . instructions . append ( get_instruction_with_outputs ( read = rsi ) ) <EOL> mt . instructions . append ( get_instruction_with_outputs ( flatten = fi ) ) <EOL> wi = dataflow . WorkItem ( ) <EOL> wi . id = <NUM_LIT> <EOL> wi . projectId = '<STR_LIT>' <EOL> wi . jobId = '<STR_LIT>' <EOL> wi . mapTask = mt <EOL> m = dataflow . LeaseWorkItemResponse ( ) <EOL> m . workItems . append ( wi ) <EOL> return m <EOL> class WorkItemTest ( unittest . TestCase ) : <EOL> def test_concat_source_to_shuffle_sink ( self ) : <EOL> work = workitem . get_work_items ( get_concat_source_to_shuffle_sink_message ( ) ) <EOL> self . assertIsNotNone ( work ) <EOL> expected_sub_sources = [ ] <EOL> expected_sub_sources . append ( <EOL> io . TextFileSource ( <EOL> file_path = '<STR_LIT>' <EOL> '<STR_LIT>' , <EOL> start_offset = <NUM_LIT:0> , end_offset = <NUM_LIT> , <EOL> strip_trailing_newlines = True , coder = CODER ) ) <EOL> expected_sub_sources . append ( <EOL> io . TextFileSource ( <EOL> file_path = '<STR_LIT>' <EOL> '<STR_LIT>' , <EOL> start_offset = <NUM_LIT:0> , end_offset = <NUM_LIT> , <EOL> strip_trailing_newlines = True , coder = CODER ) ) <EOL> expected_concat_source = concat_reader . ConcatSource ( expected_sub_sources ) <EOL> self . assertEqual ( <EOL> ( work . proto . id , work . map_task . operations ) , <EOL> ( <NUM_LIT> , [ <EOL> maptask . WorkerRead ( <EOL> expected_concat_source , output_coders = [ CODER ] ) , <EOL> maptask . WorkerDoFn ( <EOL> serialized_fn = '<STR_LIT:code>' , output_tags = [ '<STR_LIT>' ] , input = ( <NUM_LIT:1> , <NUM_LIT:0> ) , <EOL> side_inputs = [ ] , output_coders = [ CODER ] ) , <EOL> maptask . WorkerShuffleWrite ( <EOL> shuffle_kind = '<STR_LIT>' , <EOL> shuffle_writer_config = '<STR_LIT>' , <EOL> input = ( <NUM_LIT:1> , <NUM_LIT:0> ) , <EOL> output_coders = ( CODER , ) ) ] ) ) <EOL> def test_text_source_to_shuffle_sink ( self ) : <EOL> work = workitem . get_work_items ( get_text_source_to_shuffle_sink_message ( ) ) <EOL> self . assertEqual ( <EOL> ( work . proto . id , work . map_task . operations ) , <EOL> ( <NUM_LIT> , [ <EOL> maptask . WorkerRead ( io . TextFileSource ( <EOL> file_path = '<STR_LIT>' , <EOL> start_offset = <NUM_LIT> , <EOL> end_offset = <NUM_LIT> , <EOL> strip_trailing_newlines = True , <EOL> coder = CODER ) , output_coders = [ CODER ] ) , <EOL> maptask . WorkerDoFn ( <EOL> serialized_fn = '<STR_LIT:code>' , output_tags = [ '<STR_LIT>' ] , input = ( <NUM_LIT:1> , <NUM_LIT:0> ) , <EOL> side_inputs = [ ] , output_coders = [ CODER ] ) , <EOL> maptask . WorkerShuffleWrite ( <EOL> shuffle_kind = '<STR_LIT>' , <EOL> shuffle_writer_config = '<STR_LIT>' , <EOL> input = ( <NUM_LIT:1> , <NUM_LIT:0> ) , <EOL> output_coders = ( CODER , ) ) ] ) ) <EOL> def test_shuffle_source_to_text_sink ( self ) : <EOL> work = workitem . get_work_items ( <EOL> get_shuffle_source_to_text_sink_message ( GROUPING_SHUFFLE_SOURCE_SPEC ) ) <EOL> self . assertEqual ( <EOL> ( work . proto . id , work . map_task . operations ) , <EOL> ( <NUM_LIT> , [ <EOL> maptask . WorkerGroupingShuffleRead ( <EOL> start_shuffle_position = '<STR_LIT>' , <EOL> end_shuffle_position = '<STR_LIT>' , <EOL> shuffle_reader_config = '<STR_LIT>' , <EOL> coder = CODER , <EOL> output_coders = [ CODER ] ) , <EOL> maptask . WorkerWrite ( io . TextFileSink ( <EOL> file_path_prefix = '<STR_LIT>' , <EOL> append_trailing_newlines = True , <EOL> coder = CODER ) , input = ( <NUM_LIT:0> , <NUM_LIT:0> ) , output_coders = ( CODER , ) ) ] ) ) <EOL> def test_ungrouped_shuffle_source_to_text_sink ( self ) : <EOL> work = workitem . get_work_items ( <EOL> get_shuffle_source_to_text_sink_message ( UNGROUPED_SHUFFLE_SOURCE_SPEC ) ) <EOL> self . assertEqual ( <EOL> ( work . proto . id , work . map_task . operations ) , <EOL> ( <NUM_LIT> , [ <EOL> maptask . WorkerUngroupedShuffleRead ( <EOL> start_shuffle_position = '<STR_LIT>' , <EOL> end_shuffle_position = '<STR_LIT>' , <EOL> shuffle_reader_config = '<STR_LIT>' , <EOL> coder = CODER , <EOL> output_coders = [ CODER ] ) , <EOL> maptask . WorkerWrite ( io . TextFileSink ( <EOL> file_path_prefix = '<STR_LIT>' , <EOL> append_trailing_newlines = True , <EOL> coder = CODER ) , input = ( <NUM_LIT:0> , <NUM_LIT:0> ) , output_coders = ( CODER , ) ) ] ) ) <EOL> def test_in_memory_source_to_text_sink ( self ) : <EOL> work = workitem . get_work_items ( get_in_memory_source_to_text_sink_message ( ) ) <EOL> self . assertEqual ( <EOL> ( work . proto . id , work . map_task . operations ) , <EOL> ( <NUM_LIT> , [ <EOL> maptask . WorkerRead ( <EOL> inmemory . InMemorySource ( <EOL> start_index = <NUM_LIT:1> , <EOL> end_index = <NUM_LIT:3> , <EOL> elements = [ base64 . b64decode ( v [ '<STR_LIT:value>' ] ) <EOL> for v in IN_MEMORY_ELEMENTS ] , <EOL> coder = CODER ) , <EOL> output_coders = [ CODER ] ) , <EOL> maptask . WorkerWrite ( io . TextFileSink ( <EOL> file_path_prefix = '<STR_LIT>' , <EOL> append_trailing_newlines = True , <EOL> coder = CODER ) , input = ( <NUM_LIT:0> , <NUM_LIT:0> ) , output_coders = ( CODER , ) ) ] ) ) <EOL> def test_in_memory_source_to_flatten ( self ) : <EOL> work = workitem . get_work_items ( get_in_memory_source_to_flatten_message ( ) ) <EOL> self . assertEqual ( <EOL> ( work . proto . id , work . map_task . operations ) , <EOL> ( <NUM_LIT> , [ <EOL> maptask . WorkerRead ( <EOL> inmemory . InMemorySource ( <EOL> start_index = <NUM_LIT:1> , <EOL> end_index = <NUM_LIT:3> , <EOL> elements = [ base64 . b64decode ( v [ '<STR_LIT:value>' ] ) <EOL> for v in IN_MEMORY_ELEMENTS ] , <EOL> coder = CODER ) , <EOL> output_coders = [ CODER ] ) , <EOL> maptask . WorkerFlatten ( <EOL> inputs = [ ( <NUM_LIT:0> , <NUM_LIT:0> ) ] , output_coders = [ CODER ] ) ] ) ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> logging . getLogger ( ) . setLevel ( logging . INFO ) <EOL> unittest . main ( ) </s>
