Add atomically_replaced_file for safer output

[bup.git] / lib / bup / t / thelpers.py

diff --git a/lib/bup/t/thelpers.py b/lib/bup/t/thelpers.py
index b5427076dddc9e93ffd882af8a4ef9aa744408e0..02e89efb10c7b6f1c309441e9dcc46cbe408119d 100644 (file)
--- a/lib/bup/t/thelpers.py
+++ b/lib/bup/t/thelpers.py
@@ -1,6 +1,9 @@
 import helpers
 import math
 import os
+import os.path
+import tempfile
+import stat
 import bup._helpers as _helpers
 from bup.helpers import *
 from wvtest import *
@@ -143,3 +146,21 @@ def test_batchpipe():
     WVPASSEQ(next(batches), '0 1 2\n')
     WVPASSEQ(next(batches), '3 4\n')
     WVPASSEQ(next(batches, None), None)
+
+
+@wvtest
+def test_atomically_replaced_file():
+    target_file = os.path.join(tempfile.gettempdir(),
+                               'test_atomic_write')
+    try:
+        with atomically_replaced_file(target_file, mode='w') as f:
+            f.write('asdf')
+            WVPASSEQ(f.mode, 'w')
+        f = open(target_file, 'r')
+        WVPASSEQ(f.read(), 'asdf')
+
+        with atomically_replaced_file(target_file, mode='wb') as f:
+            f.write(os.urandom(20))
+            WVPASSEQ(f.mode, 'wb')
+    finally:
+        unlink(target_file)