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