@@ -789,31 +789,70 @@ def test_complete_multi_phase_init_module(self):
789789 self .run_with_own_gil (script )
790790
791791
792- class MiscTests (unittest .TestCase ):
793- def test_atomic_write_should_notice_incomplete_writes (self ):
792+ class PatchAtomicWrites :
793+ def __init__ (self , truncate_at_length , never_complete = False ):
794+ self .truncate_at_length = truncate_at_length
795+ self .never_complete = never_complete
796+ self .seen_write = False
797+ self ._children = []
798+
799+ def __enter__ (self ):
794800 import _pyio
795801
796802 oldwrite = os .write
797- seen_write = False
798-
799- truncate_at_length = 100
800803
801804 # Emulate an os.write that only writes partial data.
802805 def write (fd , data ):
803- nonlocal seen_write
804- seen_write = True
805- return oldwrite (fd , data [:truncate_at_length ])
806+ if self .seen_write and self .never_complete :
807+ return None
808+ self .seen_write = True
809+ return oldwrite (fd , data [:self .truncate_at_length ])
806810
807811 # Need to patch _io to be _pyio, so that io.FileIO is affected by the
808812 # os.write patch.
809- with (support .swap_attr (_bootstrap_external , '_io' , _pyio ),
810- support .swap_attr (os , 'write' , write )):
811- with self .assertRaises (OSError ):
812- # Make sure we write something longer than the point where we
813- # truncate.
814- content = b'x' * (truncate_at_length * 2 )
815- _bootstrap_external ._write_atomic (os_helper .TESTFN , content )
816- assert seen_write
813+ self .children = [
814+ support .swap_attr (_bootstrap_external , '_io' , _pyio ),
815+ support .swap_attr (os , 'write' , write )
816+ ]
817+ for child in self .children :
818+ child .__enter__ ()
819+ return self
820+
821+ def __exit__ (self , exc_type , exc_val , exc_tb ):
822+ for child in self .children :
823+ child .__exit__ (exc_type , exc_val , exc_tb )
824+
825+
826+ class MiscTests (unittest .TestCase ):
827+
828+ def test_atomic_write_retries_incomplete_writes (self ):
829+ truncate_at_length = 100
830+ length = truncate_at_length * 2
831+
832+ with PatchAtomicWrites (truncate_at_length = truncate_at_length ) as cm :
833+ # Make sure we write something longer than the point where we
834+ # truncate.
835+ content = b'x' * length
836+ _bootstrap_external ._write_atomic (os_helper .TESTFN , content )
837+ self .assertTrue (cm .seen_write )
838+
839+ self .assertEqual (os .stat (support .os_helper .TESTFN ).st_size , length )
840+ os .unlink (support .os_helper .TESTFN )
841+
842+ def test_atomic_write_errors_if_unable_to_complete (self ):
843+ truncate_at_length = 100
844+
845+ with (
846+ PatchAtomicWrites (
847+ truncate_at_length = truncate_at_length , never_complete = True ,
848+ ) as cm ,
849+ self .assertRaises (OSError )
850+ ):
851+ # Make sure we write something longer than the point where we
852+ # truncate.
853+ content = b'x' * (truncate_at_length * 2 )
854+ _bootstrap_external ._write_atomic (os_helper .TESTFN , content )
855+ self .assertTrue (cm .seen_write )
817856
818857 with self .assertRaises (OSError ):
819858 os .stat (support .os_helper .TESTFN ) # Check that the file did not get written.
0 commit comments