<s> import unittest <EOL> import pipeline <EOL> import cloudstorage <EOL> from mapreduce import base_handler <EOL> from mapreduce import kv_pb <EOL> from mapreduce import mapreduce_pipeline <EOL> from mapreduce import output_writers <EOL> from mapreduce import records <EOL> from mapreduce import shuffler <EOL> from mapreduce import test_support <EOL> from testlib import testutil <EOL> class HashEndToEndTest ( testutil . HandlerTestBase ) : <EOL> """<STR_LIT>""" <EOL> def setUp ( self ) : <EOL> testutil . HandlerTestBase . setUp ( self ) <EOL> pipeline . Pipeline . _send_mail = self . _send_mail <EOL> self . emails = [ ] <EOL> def _send_mail ( self , sender , subject , body , html = None ) : <EOL> """<STR_LIT>""" <EOL> self . emails . append ( ( sender , subject , body , html ) ) <EOL> def testHashingMultipleFiles ( self ) : <EOL> """<STR_LIT>""" <EOL> input_data = [ ( str ( i ) , str ( i ) ) for i in range ( <NUM_LIT:100> ) ] <EOL> input_data . sort ( ) <EOL> bucket_name = "<STR_LIT>" <EOL> test_filename = "<STR_LIT>" <EOL> full_filename = "<STR_LIT>" % ( bucket_name , test_filename ) <EOL> with cloudstorage . open ( full_filename , mode = "<STR_LIT:w>" ) as f : <EOL> with records . RecordsWriter ( f ) as w : <EOL> for ( k , v ) in input_data : <EOL> proto = kv_pb . KeyValue ( ) <EOL> proto . set_key ( k ) <EOL> proto . set_value ( v ) <EOL> w . write ( proto . Encode ( ) ) <EOL> p = shuffler . _HashPipeline ( "<STR_LIT>" , bucket_name , <EOL> [ full_filename , full_filename , full_filename ] ) <EOL> p . start ( ) <EOL> test_support . execute_until_empty ( self . taskqueue ) <EOL> p = shuffler . _HashPipeline . from_id ( p . pipeline_id ) <EOL> list_of_output_files = p . outputs . default . value <EOL> output_data = [ ] <EOL> for output_files in list_of_output_files : <EOL> for output_file in output_files : <EOL> with cloudstorage . open ( output_file ) as f : <EOL> for binary_record in records . RecordsReader ( f ) : <EOL> proto = kv_pb . KeyValue ( ) <EOL> proto . ParseFromString ( binary_record ) <EOL> output_data . append ( ( proto . key ( ) , proto . value ( ) ) ) <EOL> output_data . sort ( ) <EOL> self . assertEquals ( <NUM_LIT> , len ( output_data ) ) <EOL> for i in range ( len ( input_data ) ) : <EOL> self . assertEquals ( input_data [ i ] , output_data [ ( <NUM_LIT:3> * i ) ] ) <EOL> self . assertEquals ( input_data [ i ] , output_data [ ( <NUM_LIT:3> * i ) + <NUM_LIT:1> ] ) <EOL> self . assertEquals ( input_data [ i ] , output_data [ ( <NUM_LIT:3> * i ) + <NUM_LIT:2> ] ) <EOL> self . assertEquals ( <NUM_LIT:1> , len ( self . emails ) ) <EOL> class SortFileEndToEndTest ( testutil . HandlerTestBase ) : <EOL> """<STR_LIT>""" <EOL> def setUp ( self ) : <EOL> testutil . HandlerTestBase . setUp ( self ) <EOL> pipeline . Pipeline . _send_mail = self . _send_mail <EOL> self . emails = [ ] <EOL> def _send_mail ( self , sender , subject , body , html = None ) : <EOL> """<STR_LIT>""" <EOL> self . emails . append ( ( sender , subject , body , html ) ) <EOL> def testSortFile ( self ) : <EOL> """<STR_LIT>""" <EOL> bucket_name = "<STR_LIT>" <EOL> test_filename = "<STR_LIT>" <EOL> full_filename = "<STR_LIT>" % ( bucket_name , test_filename ) <EOL> input_data = [ <EOL> ( str ( i ) , "<STR_LIT:_>" + str ( i ) ) for i in range ( <NUM_LIT:100> ) ] <EOL> with cloudstorage . open ( full_filename , mode = "<STR_LIT:w>" ) as f : <EOL> with records . RecordsWriter ( f ) as w : <EOL> for ( k , v ) in input_data : <EOL> proto = kv_pb . KeyValue ( ) <EOL> proto . set_key ( k ) <EOL> proto . set_value ( v ) <EOL> w . write ( proto . Encode ( ) ) <EOL> p = shuffler . _SortChunksPipeline ( "<STR_LIT>" , bucket_name , [ [ full_filename ] ] ) <EOL> p . start ( ) <EOL> test_support . execute_until_empty ( self . taskqueue ) <EOL> p = shuffler . _SortChunksPipeline . from_id ( p . pipeline_id ) <EOL> input_data . sort ( ) <EOL> output_files = p . outputs . default . value [ <NUM_LIT:0> ] <EOL> output_data = [ ] <EOL> for output_file in output_files : <EOL> with cloudstorage . open ( output_file ) as f : <EOL> for binary_record in records . RecordsReader ( f ) : <EOL> proto = kv_pb . KeyValue ( ) <EOL> proto . ParseFromString ( binary_record ) <EOL> output_data . append ( ( proto . key ( ) , proto . value ( ) ) ) <EOL> self . assertEquals ( input_data , output_data ) <EOL> self . assertEquals ( <NUM_LIT:1> , len ( self . emails ) ) <EOL> def test_handler_yield_str ( key , value , partial ) : <EOL> """<STR_LIT>""" <EOL> yield str ( ( key , value , partial ) ) <EOL> class TestMergePipeline ( base_handler . PipelineBase ) : <EOL> """<STR_LIT>""" <EOL> def run ( self , bucket_name , filenames ) : <EOL> yield mapreduce_pipeline . MapperPipeline ( <EOL> "<STR_LIT>" , <EOL> __name__ + "<STR_LIT>" , <EOL> shuffler . __name__ + "<STR_LIT>" , <EOL> output_writers . __name__ + "<STR_LIT>" , <EOL> params = { <EOL> shuffler . _MergingReader . FILES_PARAM : [ filenames ] , <EOL> shuffler . _MergingReader . MAX_VALUES_COUNT_PARAM : <EOL> shuffler . _MergePipeline . _MAX_VALUES_COUNT , <EOL> shuffler . _MergingReader . MAX_VALUES_SIZE_PARAM : <EOL> shuffler . _MergePipeline . _MAX_VALUES_SIZE , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : bucket_name , <EOL> } , <EOL> } , <EOL> ) <EOL> class MergingReaderEndToEndTest ( testutil . HandlerTestBase ) : <EOL> """<STR_LIT>""" <EOL> def setUp ( self ) : <EOL> testutil . HandlerTestBase . setUp ( self ) <EOL> pipeline . Pipeline . _send_mail = self . _send_mail <EOL> self . emails = [ ] <EOL> def _send_mail ( self , sender , subject , body , html = None ) : <EOL> """<STR_LIT>""" <EOL> self . emails . append ( ( sender , subject , body , html ) ) <EOL> def testMergeFiles ( self ) : <EOL> """<STR_LIT>""" <EOL> input_data = [ ( str ( i ) , "<STR_LIT:_>" + str ( i ) ) for i in range ( <NUM_LIT:100> ) ] <EOL> input_data . sort ( ) <EOL> bucket_name = "<STR_LIT>" <EOL> test_filename = "<STR_LIT>" <EOL> full_filename = "<STR_LIT>" % ( bucket_name , test_filename ) <EOL> with cloudstorage . open ( full_filename , mode = "<STR_LIT:w>" ) as f : <EOL> with records . RecordsWriter ( f ) as w : <EOL> for ( k , v ) in input_data : <EOL> proto = kv_pb . KeyValue ( ) <EOL> proto . set_key ( k ) <EOL> proto . set_value ( v ) <EOL> w . write ( proto . Encode ( ) ) <EOL> p = TestMergePipeline ( bucket_name , <EOL> [ full_filename , full_filename , full_filename ] ) <EOL> p . start ( ) <EOL> test_support . execute_until_empty ( self . taskqueue ) <EOL> p = TestMergePipeline . from_id ( p . pipeline_id ) <EOL> output_file = p . outputs . default . value [ <NUM_LIT:0> ] <EOL> output_data = [ ] <EOL> with cloudstorage . open ( output_file ) as f : <EOL> for record in records . RecordsReader ( f ) : <EOL> output_data . append ( record ) <EOL> expected_data = [ <EOL> str ( ( k , [ v , v , v ] , False ) ) for ( k , v ) in input_data ] <EOL> self . assertEquals ( expected_data , output_data ) <EOL> self . assertEquals ( <NUM_LIT:1> , len ( self . emails ) ) <EOL> def testPartialRecords ( self ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> self . _prev_max_values_count = shuffler . _MergePipeline . _MAX_VALUES_COUNT <EOL> shuffler . _MergePipeline . _MAX_VALUES_COUNT = <NUM_LIT:1> <EOL> input_data = [ ( "<STR_LIT:1>" , "<STR_LIT:a>" ) , ( "<STR_LIT:2>" , "<STR_LIT:b>" ) , ( "<STR_LIT:3>" , "<STR_LIT:c>" ) ] <EOL> input_data . sort ( ) <EOL> bucket_name = "<STR_LIT>" <EOL> test_filename = "<STR_LIT>" <EOL> full_filename = "<STR_LIT>" % ( bucket_name , test_filename ) <EOL> with cloudstorage . open ( full_filename , mode = "<STR_LIT:w>" ) as f : <EOL> with records . RecordsWriter ( f ) as w : <EOL> for ( k , v ) in input_data : <EOL> proto = kv_pb . KeyValue ( ) <EOL> proto . set_key ( k ) <EOL> proto . set_value ( v ) <EOL> w . write ( proto . Encode ( ) ) <EOL> p = TestMergePipeline ( bucket_name , <EOL> [ full_filename , full_filename , full_filename ] ) <EOL> p . start ( ) <EOL> test_support . execute_until_empty ( self . taskqueue ) <EOL> p = TestMergePipeline . from_id ( p . pipeline_id ) <EOL> output_file = p . outputs . default . value [ <NUM_LIT:0> ] <EOL> output_data = [ ] <EOL> with cloudstorage . open ( output_file ) as f : <EOL> for record in records . RecordsReader ( f ) : <EOL> output_data . append ( record ) <EOL> expected_data = [ <EOL> ( "<STR_LIT:1>" , [ "<STR_LIT:a>" ] , True ) , <EOL> ( "<STR_LIT:1>" , [ "<STR_LIT:a>" ] , True ) , <EOL> ( "<STR_LIT:1>" , [ "<STR_LIT:a>" ] , False ) , <EOL> ( "<STR_LIT:2>" , [ "<STR_LIT:b>" ] , True ) , <EOL> ( "<STR_LIT:2>" , [ "<STR_LIT:b>" ] , True ) , <EOL> ( "<STR_LIT:2>" , [ "<STR_LIT:b>" ] , False ) , <EOL> ( "<STR_LIT:3>" , [ "<STR_LIT:c>" ] , True ) , <EOL> ( "<STR_LIT:3>" , [ "<STR_LIT:c>" ] , True ) , <EOL> ( "<STR_LIT:3>" , [ "<STR_LIT:c>" ] , False ) , <EOL> ] <EOL> self . assertEquals ( [ str ( e ) for e in expected_data ] , output_data ) <EOL> finally : <EOL> shuffler . _MergePipeline . _MAX_VALUES_COUNT = self . _prev_max_values_count <EOL> self . assertEquals ( <NUM_LIT:1> , len ( self . emails ) ) <EOL> class ShuffleEndToEndTest ( testutil . HandlerTestBase ) : <EOL> """<STR_LIT>""" <EOL> def setUp ( self ) : <EOL> testutil . HandlerTestBase . setUp ( self ) <EOL> pipeline . Pipeline . _send_mail = self . _send_mail <EOL> self . emails = [ ] <EOL> def _send_mail ( self , sender , subject , body , html = None ) : <EOL> """<STR_LIT>""" <EOL> self . emails . append ( ( sender , subject , body , html ) ) <EOL> def testShuffleNoData ( self ) : <EOL> bucket_name = "<STR_LIT>" <EOL> test_filename = "<STR_LIT>" <EOL> full_filename = "<STR_LIT>" % ( bucket_name , test_filename ) <EOL> gcs_file = cloudstorage . open ( full_filename , mode = "<STR_LIT:w>" ) <EOL> gcs_file . close ( ) <EOL> p = shuffler . ShufflePipeline ( "<STR_LIT>" , { "<STR_LIT>" : bucket_name } , <EOL> [ full_filename , full_filename , full_filename ] ) <EOL> p . start ( ) <EOL> test_support . execute_until_empty ( self . taskqueue ) <EOL> p = shuffler . ShufflePipeline . from_id ( p . pipeline_id ) <EOL> for filename in p . outputs . default . value : <EOL> self . assertEqual ( <NUM_LIT:0> , cloudstorage . stat ( filename ) . st_size ) <EOL> self . assertEquals ( <NUM_LIT:1> , len ( self . emails ) ) <EOL> def testShuffleNoFile ( self ) : <EOL> bucket_name = "<STR_LIT>" <EOL> p = shuffler . ShufflePipeline ( "<STR_LIT>" , { "<STR_LIT>" : bucket_name } , [ ] ) <EOL> p . start ( ) <EOL> test_support . execute_until_empty ( self . taskqueue ) <EOL> p = shuffler . ShufflePipeline . from_id ( p . pipeline_id ) <EOL> for filename in p . outputs . default . value : <EOL> self . assertEqual ( <NUM_LIT:0> , cloudstorage . stat ( filename ) . st_size ) <EOL> self . assertEquals ( <NUM_LIT:1> , len ( self . emails ) ) <EOL> def testShuffleFiles ( self ) : <EOL> """<STR_LIT>""" <EOL> input_data = [ ( str ( i ) , str ( i ) ) for i in range ( <NUM_LIT:100> ) ] <EOL> input_data . sort ( ) <EOL> bucket_name = "<STR_LIT>" <EOL> test_filename = "<STR_LIT>" <EOL> full_filename = "<STR_LIT>" % ( bucket_name , test_filename ) <EOL> with cloudstorage . open ( full_filename , mode = "<STR_LIT:w>" ) as f : <EOL> with records . RecordsWriter ( f ) as w : <EOL> for ( k , v ) in input_data : <EOL> proto = kv_pb . KeyValue ( ) <EOL> proto . set_key ( k ) <EOL> proto . set_value ( v ) <EOL> w . write ( proto . Encode ( ) ) <EOL> p = shuffler . ShufflePipeline ( "<STR_LIT>" , { "<STR_LIT>" : bucket_name } , <EOL> [ full_filename , full_filename , full_filename ] ) <EOL> p . start ( ) <EOL> test_support . execute_until_empty ( self . taskqueue ) <EOL> p = shuffler . ShufflePipeline . from_id ( p . pipeline_id ) <EOL> output_files = p . outputs . default . value <EOL> output_data = [ ] <EOL> for output_file in output_files : <EOL> with cloudstorage . open ( output_file ) as f : <EOL> for record in records . RecordsReader ( f ) : <EOL> proto = kv_pb . KeyValues ( ) <EOL> proto . ParseFromString ( record ) <EOL> output_data . append ( ( proto . key ( ) , proto . value_list ( ) ) ) <EOL> output_data . sort ( ) <EOL> expected_data = sorted ( [ <EOL> ( str ( k ) , [ str ( v ) , str ( v ) , str ( v ) ] ) for ( k , v ) in input_data ] ) <EOL> self . assertEquals ( expected_data , output_data ) <EOL> self . assertEquals ( <NUM_LIT:1> , len ( self . emails ) ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> unittest . main ( ) </s>
