<s> import gevent <EOL> import pytest <EOL> from fast_wait import fast_wait <EOL> from wal_e import worker <EOL> from wal_e . exception import UserCritical <EOL> assert fast_wait <EOL> class Explosion ( Exception ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> class FakeWalSegment ( object ) : <EOL> def __init__ ( self , seg_path , explicit = False , <EOL> upload_explosive = False , <EOL> mark_done_explosive = False ) : <EOL> self . explicit = explicit <EOL> self . _upload_explosive = upload_explosive <EOL> self . _mark_done_explosive = mark_done_explosive <EOL> self . _marked = False <EOL> self . _uploaded = False <EOL> def mark_done ( self ) : <EOL> if self . _mark_done_explosive : <EOL> raise self . _mark_done_explosive <EOL> self . _marked = True <EOL> class FakeWalUploader ( object ) : <EOL> def __call__ ( self , segment ) : <EOL> if segment . _upload_explosive : <EOL> raise segment . _upload_explosive <EOL> segment . _uploaded = True <EOL> return segment <EOL> def failed ( seg ) : <EOL> """<STR_LIT>""" <EOL> return seg . _marked is False and seg . _uploaded is False <EOL> def success ( seg ) : <EOL> """<STR_LIT>""" <EOL> if seg . explicit : <EOL> assert seg . _marked is False <EOL> return seg . _uploaded <EOL> def indeterminate ( seg ) : <EOL> """<STR_LIT>""" <EOL> if seg . _uploaded : <EOL> if seg . explicit : <EOL> assert seg . _marked is False <EOL> else : <EOL> assert seg . _marked is True <EOL> else : <EOL> assert seg . _marked is False <EOL> return True <EOL> def prepare_multi_upload_segments ( ) : <EOL> """<STR_LIT>""" <EOL> yield FakeWalSegment ( '<STR_LIT:0>' * <NUM_LIT:8> * <NUM_LIT:3> , explicit = True ) <EOL> for i in xrange ( <NUM_LIT:1> , <NUM_LIT:5> ) : <EOL> yield FakeWalSegment ( str ( i ) * <NUM_LIT:8> * <NUM_LIT:3> , explicit = False ) <EOL> def test_simple_upload ( ) : <EOL> """<STR_LIT>""" <EOL> group = worker . WalTransferGroup ( FakeWalUploader ( ) ) <EOL> seg = FakeWalSegment ( '<STR_LIT:1>' * <NUM_LIT:8> * <NUM_LIT:3> , explicit = True ) <EOL> group . start ( seg ) <EOL> group . join ( ) <EOL> assert success ( seg ) <EOL> def test_multi_upload ( ) : <EOL> """<STR_LIT>""" <EOL> group = worker . WalTransferGroup ( FakeWalUploader ( ) ) <EOL> segments = list ( prepare_multi_upload_segments ( ) ) <EOL> for seg in segments : <EOL> group . start ( seg ) <EOL> group . join ( ) <EOL> for seg in segments : <EOL> assert success ( seg ) <EOL> def test_simple_fail ( ) : <EOL> """<STR_LIT>""" <EOL> group = worker . WalTransferGroup ( FakeWalUploader ( ) ) <EOL> exp = Explosion ( '<STR_LIT>' ) <EOL> seg = FakeWalSegment ( '<STR_LIT:1>' * <NUM_LIT:8> * <NUM_LIT:3> , explicit = True , upload_explosive = exp ) <EOL> group . start ( seg ) <EOL> with pytest . raises ( Explosion ) as e : <EOL> group . join ( ) <EOL> assert e . value is exp <EOL> assert failed ( seg ) <EOL> def test_multi_explicit_fail ( ) : <EOL> """<STR_LIT>""" <EOL> group = worker . WalTransferGroup ( FakeWalUploader ( ) ) <EOL> segments = list ( prepare_multi_upload_segments ( ) ) <EOL> exp = Explosion ( '<STR_LIT>' ) <EOL> segments [ <NUM_LIT:0> ] . _upload_explosive = exp <EOL> for seg in segments : <EOL> group . start ( seg ) <EOL> with pytest . raises ( Explosion ) as e : <EOL> group . join ( ) <EOL> assert e . value is exp <EOL> assert failed ( segments [ <NUM_LIT:0> ] ) <EOL> for seg in segments [ <NUM_LIT:1> : ] : <EOL> assert success ( seg ) <EOL> def test_multi_pipeline_fail ( ) : <EOL> """<STR_LIT>""" <EOL> group = worker . WalTransferGroup ( FakeWalUploader ( ) ) <EOL> segments = list ( prepare_multi_upload_segments ( ) ) <EOL> exp = Explosion ( '<STR_LIT>' ) <EOL> fail_idx = <NUM_LIT:2> <EOL> segments [ fail_idx ] . _upload_explosive = exp <EOL> for seg in segments : <EOL> group . start ( seg ) <EOL> with pytest . raises ( Explosion ) as e : <EOL> group . join ( ) <EOL> assert e . value is exp <EOL> for i , seg in enumerate ( segments ) : <EOL> if i == fail_idx : <EOL> assert failed ( seg ) <EOL> else : <EOL> assert indeterminate ( seg ) <EOL> def test_finally_execution ( ) : <EOL> """<STR_LIT>""" <EOL> segBad = FakeWalSegment ( '<STR_LIT:1>' * <NUM_LIT:8> * <NUM_LIT:3> ) <EOL> segOK = FakeWalSegment ( '<STR_LIT:2>' * <NUM_LIT:8> * <NUM_LIT:3> ) <EOL> class CleanupCheckingUploader ( object ) : <EOL> def __init__ ( self ) : <EOL> self . cleaned_up = False <EOL> def __call__ ( self , segment ) : <EOL> if segment is segOK : <EOL> try : <EOL> while True : <EOL> gevent . sleep ( <NUM_LIT:0.1> ) <EOL> finally : <EOL> self . cleaned_up = True <EOL> elif segment is segBad : <EOL> raise Explosion ( '<STR_LIT>' ) <EOL> else : <EOL> assert False , '<STR_LIT>' <EOL> segment . _uploaded = True <EOL> return segment <EOL> uploader = CleanupCheckingUploader ( ) <EOL> group = worker . WalTransferGroup ( uploader ) <EOL> group . start ( segOK ) <EOL> group . start ( segBad ) <EOL> with pytest . raises ( Explosion ) : <EOL> group . join ( ) <EOL> assert uploader . cleaned_up is True <EOL> def test_start_after_join ( ) : <EOL> """<STR_LIT>""" <EOL> group = worker . WalTransferGroup ( FakeWalUploader ( ) ) <EOL> group . join ( ) <EOL> seg = FakeWalSegment ( '<STR_LIT>' ) <EOL> with pytest . raises ( UserCritical ) : <EOL> group . start ( seg ) <EOL> def test_mark_done_fault ( ) : <EOL> """<STR_LIT>""" <EOL> group = worker . WalTransferGroup ( FakeWalUploader ( ) ) <EOL> exp = Explosion ( '<STR_LIT>' ) <EOL> seg = FakeWalSegment ( '<STR_LIT>' , mark_done_explosive = exp ) <EOL> group . start ( seg ) <EOL> with pytest . raises ( Explosion ) as e : <EOL> group . join ( ) <EOL> assert e . value is exp </s>
